Repository: neurogym/neurogym

Issue:
neurogym/utils/random.py shadows python standard library
The file name `random` is not great, because it is identical to a python built-in library. 

It would be good to rename this file to distinguish between it and the built-in package. This would be a breaking change, however, and I am not sure we want to immediately move to a version 2.0. 

Note that this issue is also flagged by ruff > 0.9 rule [A005](https://docs.astral.sh/ruff/rules/stdlib-module-shadowing/), which is why we are currently still using ruff 0.8.x.

Discussion:

Solution patch:
diff --git a/.ruff.toml b/.ruff.toml
index 4692ce91..5f390694 100644
--- a/.ruff.toml
+++ b/.ruff.toml
@@ -32,8 +32,10 @@ ignore = [
     "D107",   # Missing docstring in `__init__`
     "CPY001", # Missing copyright notice
     # preview ignores
-    "PLR0917", # too many positional arguments
-    "PLR0914", # too many local arguments
+    "PLR0917", # Too many positional arguments
+    "PLR0914", # Too many local arguments
+    "DOC201",  # Return is not documented in docstring
+    "DOC501",  # Exception is not documented in docstring
 ]
 
 # Autofix settings
@@ -67,10 +69,6 @@ extend-safe-fixes = [
     # "T201",   # Print statements
 ]
 "__init__.py" = ["F401"] # unused imports
-# files and locations that will likely not be preserved into new version
-"docs/*" = ["ALL"]
-"setup.py" = ["ALL"]
-
 
 [lint.flake8-unused-arguments]
 ignore-variadic-names = true # Whether to allow unused variadic arguments, like *args and **kwargs.
diff --git a/neurogym/__init__.py b/neurogym/__init__.py
index 4172cb25..9ee1503a 100644
--- a/neurogym/__init__.py
+++ b/neurogym/__init__.py
@@ -13,6 +13,6 @@ from neurogym.envs.registration import all_tags
 from neurogym.envs.collections import get_collection
 from neurogym.wrappers import all_wrappers
 from neurogym.utils.data import Dataset
-from neurogym.utils import random
+from neurogym.utils import ngym_random
 
 __version__ = "1.0.3"
diff --git a/neurogym/core.py b/neurogym/core.py
index 6a354bac..76f260ed 100644
--- a/neurogym/core.py
+++ b/neurogym/core.py
@@ -4,7 +4,7 @@ from typing import Any, NoReturn
 import gymnasium as gym
 import numpy as np
 
-from neurogym.utils.random import trunc_exp
+from neurogym.utils.ngym_random import trunc_exp
 
 METADATA_DEF_KEYS = ["description", "paper_name", "paper_link", "timing", "tags"]
 
@@ -140,7 +140,8 @@ class TrialEnv(BaseEnv):
                 val.seed(seed)
         return [seed]
 
-    def post_step(self, ob, reward, terminated, truncated, info):
+    @staticmethod
+    def post_step(ob, reward, terminated, truncated, info):
         """Optional task-specific wrapper applied at the end of step.
 
         It allows to modify ob online (e.g. provide a specific observation for different actions made by the agent)
diff --git a/neurogym/envs/bandit.py b/neurogym/envs/bandit.py
index e52eebe9..614e19bb 100644
--- a/neurogym/envs/bandit.py
+++ b/neurogym/envs/bandit.py
@@ -31,8 +31,8 @@ class Bandit(ngym.TrialEnv):
         dt: int = 100,
         n: int = 2,
         p: tuple[float, ...] | list[float] = (0.5, 0.5),
-        rewards: None | list[float] | np.ndarray = None,
-        timing: None | dict = None,
+        rewards: list[float] | np.ndarray | None = None,
+        timing: dict | None = None,
     ) -> None:
         super().__init__(dt=dt)
         if timing is not None:
diff --git a/neurogym/envs/contextdecisionmaking.py b/neurogym/envs/contextdecisionmaking.py
index 2ab51066..78231aa7 100644
--- a/neurogym/envs/contextdecisionmaking.py
+++ b/neurogym/envs/contextdecisionmaking.py
@@ -2,6 +2,7 @@ import numpy as np
 
 import neurogym as ngym
 from neurogym import spaces
+from neurogym.utils.ngym_random import TruncExp
 
 
 class SingleContextDecisionMaking(ngym.TrialEnv):
@@ -52,7 +53,7 @@ class SingleContextDecisionMaking(ngym.TrialEnv):
             "fixation": 300,
             # 'target': 350, # noqa: ERA001
             "stimulus": 750,
-            "delay": ngym.random.TruncExp(600, 300, 3000),
+            "delay": TruncExp(600, 300, 3000),
             "decision": 100,
         }
         if timing:
@@ -173,7 +174,7 @@ class ContextDecisionMaking(ngym.TrialEnv):
             "fixation": 300,
             # 'target': 350, # noqa: ERA001
             "stimulus": 750,
-            "delay": ngym.random.TruncExp(600, 300, 3000),
+            "delay": TruncExp(600, 300, 3000),
             "decision": 100,
         }
         if timing:
diff --git a/neurogym/envs/contrib/changingenvironment.py b/neurogym/envs/contrib/changingenvironment.py
index e82e3e8c..7acdc1ce 100644
--- a/neurogym/envs/contrib/changingenvironment.py
+++ b/neurogym/envs/contrib/changingenvironment.py
@@ -7,6 +7,7 @@ import numpy as np
 from gymnasium import spaces
 
 import neurogym as ngym
+from neurogym.utils.ngym_random import TruncExp
 
 
 # TODO: Need a more intuitive name
@@ -57,7 +58,7 @@ class ChangingEnvironment(ngym.TrialEnv):
 
         self.timing = {
             "fixation": 500,
-            "stimulus": ngym.random.TruncExp(1000, 500, 1500),
+            "stimulus": TruncExp(1000, 500, 1500),
             "decision": 500,
         }
         if timing:
diff --git a/neurogym/envs/contrib/leverpress.py b/neurogym/envs/contrib/leverpress.py
index 6cf0c9a1..6143c7ba 100644
--- a/neurogym/envs/contrib/leverpress.py
+++ b/neurogym/envs/contrib/leverpress.py
@@ -242,9 +242,6 @@ class LeverPressWithPokeRest(gym.Env):
         self.np_random, seed = seeding.np_random(seed)
         return [seed]
 
-    def _get_thirst(self, thirst_state):
-        return (thirst_state > 0.0) * (thirst_state < 1.0) * thirst_state + (thirst_state > 1.0) * 1.0
-
     def _step(self, action):
         if not self.action_space.contains(action):
             msg = f"{action=!r} ({type(action)}) not found in action_space."
@@ -272,7 +269,7 @@ class LeverPressWithPokeRest(gym.Env):
             raise ValueError
 
         self.thirst_state += self.np_random.rand() * 0.4 + 0.8
-        self.thirst = self._get_thirst(self.thirst_state)
+        self.thirst = np.clip(self.thirst_state, 0.0, 1.0)
         self.state = state
         terminated = False
         truncated = False
@@ -311,7 +308,7 @@ class LeverPressWithPokeRest(gym.Env):
         self.state = self.n_press
         self.steps_beyond_done = None
         self.thirst_state = 1
-        self.thirst = self._get_thirst(self.thirst_state)
+        self.thirst = np.clip(self.thirst_state, 0.0, 1.0)
         if self.observe_state:
             return np.array([self.state]), {}
         return np.array([self.thirst]), {}
diff --git a/neurogym/envs/contrib/pneumostomeopening.py b/neurogym/envs/contrib/pneumostomeopening.py
index c14e3be0..c40d4afc 100755
--- a/neurogym/envs/contrib/pneumostomeopening.py
+++ b/neurogym/envs/contrib/pneumostomeopening.py
@@ -64,7 +64,7 @@ class Pneumostomeopening(ngym.BaseEnv):
             self.oxygen_level += 1
         else:  # swimming actions
             self.oxygen_level *= math.exp(-decay_constant * self.t)
-        self.oxygen_level = int(round(self.oxygen_level))
+        self.oxygen_level = round(self.oxygen_level)
         print(f"oxygen_level: {self.oxygen_level}")
         return self.oxygen_level
 
diff --git a/neurogym/envs/detection.py b/neurogym/envs/detection.py
index 021358c4..725f0d9e 100644
--- a/neurogym/envs/detection.py
+++ b/neurogym/envs/detection.py
@@ -9,6 +9,7 @@ import numpy as np
 
 import neurogym as ngym
 from neurogym import spaces
+from neurogym.utils.ngym_random import TruncExp
 
 
 class Detection(ngym.TrialEnv):
@@ -67,7 +68,7 @@ class Detection(ngym.TrialEnv):
 
         self.timing = {
             "fixation": 500,
-            "stimulus": ngym.random.TruncExp(1000, 500, 1500),
+            "stimulus": TruncExp(1000, 500, 1500),
         }
         if timing:
             self.timing.update(timing)
diff --git a/neurogym/envs/hierarchicalreasoning.py b/neurogym/envs/hierarchicalreasoning.py
index ef8f75b3..e6cd2976 100644
--- a/neurogym/envs/hierarchicalreasoning.py
+++ b/neurogym/envs/hierarchicalreasoning.py
@@ -4,6 +4,7 @@ import numpy as np
 
 import neurogym as ngym
 from neurogym import spaces
+from neurogym.utils.ngym_random import TruncExp
 
 
 class HierarchicalReasoning(ngym.TrialEnv):
@@ -35,9 +36,9 @@ class HierarchicalReasoning(ngym.TrialEnv):
             self.rewards.update(rewards)
 
         self.timing = {
-            "fixation": ngym.random.TruncExp(600, 400, 800),
+            "fixation": TruncExp(600, 400, 800),
             "rule_target": 1000,
-            "fixation2": ngym.random.TruncExp(600, 400, 900),
+            "fixation2": TruncExp(600, 400, 900),
             "flash1": 100,
             "delay": (530, 610, 690, 770, 850, 930, 1010, 1090, 1170),
             "flash2": 100,
diff --git a/neurogym/envs/postdecisionwager.py b/neurogym/envs/postdecisionwager.py
index 6b04ae5a..0e553558 100644
--- a/neurogym/envs/postdecisionwager.py
+++ b/neurogym/envs/postdecisionwager.py
@@ -2,6 +2,7 @@ import numpy as np
 
 import neurogym as ngym
 from neurogym import spaces
+from neurogym.utils.ngym_random import TruncExp
 
 
 class PostDecisionWager(ngym.TrialEnv):
@@ -40,8 +41,8 @@ class PostDecisionWager(ngym.TrialEnv):
         self.timing = {
             "fixation": 100,
             # 'target':  0,  # noqa: ERA001
-            "stimulus": ngym.random.TruncExp(180, 100, 900),
-            "delay": ngym.random.TruncExp(1350, 1200, 1800),
+            "stimulus": TruncExp(180, 100, 900),
+            "delay": TruncExp(1350, 1200, 1800),
             "pre_sure": lambda: self.rng.uniform(500, 750),
             "decision": 100,
         }
@@ -63,7 +64,8 @@ class PostDecisionWager(ngym.TrialEnv):
         self.action_space = spaces.Discrete(4, name=name)
 
     # Input scaling
-    def scale(self, coh):
+    @staticmethod
+    def scale(coh):
         return (1 + coh / 100) / 2
 
     def _new_trial(self, **kwargs):
diff --git a/neurogym/envs/readysetgo.py b/neurogym/envs/readysetgo.py
index bc961e30..3afd215f 100644
--- a/neurogym/envs/readysetgo.py
+++ b/neurogym/envs/readysetgo.py
@@ -4,6 +4,7 @@ import numpy as np
 
 import neurogym as ngym
 from neurogym import spaces
+from neurogym.utils.ngym_random import TruncExp
 
 
 class ReadySetGo(ngym.TrialEnv):
@@ -242,8 +243,8 @@ class OneTwoThreeGo(ngym.TrialEnv):
             self.rewards.update(rewards)
 
         self.timing = {
-            "fixation": ngym.random.TruncExp(400, 100, 800),
-            "target": ngym.random.TruncExp(1000, 500, 1500),
+            "fixation": TruncExp(400, 100, 800),
+            "target": TruncExp(1000, 500, 1500),
             "s1": 100,
             "interval1": (600, 700, 800, 900, 1000),
             "s2": 100,
diff --git a/neurogym/utils/random.py b/neurogym/utils/ngym_random.py
similarity index 100%
rename from neurogym/utils/random.py
rename to neurogym/utils/ngym_random.py
diff --git a/neurogym/utils/plotting.py b/neurogym/utils/plotting.py
index 8e3f5404..22fb9f5c 100644
--- a/neurogym/utils/plotting.py
+++ b/neurogym/utils/plotting.py
@@ -104,11 +104,11 @@ def run_env(env, num_steps=200, num_trials=None, def_act=None, model=None):
     trial_count = 0
     for _ in range(int(num_steps)):
         if model is not None:
-            action, _states = model.predict(ob)
+            action, states = model.predict(ob)
             if isinstance(action, float | int):
                 action = [action]
-            if (_states is not None) and (len(_states) > 0):
-                state_mat.append(_states)
+            if (states is not None) and (len(states) > 0):
+                state_mat.append(states)
         elif def_act is not None:
             action = def_act
         else:
diff --git a/pyproject.toml b/pyproject.toml
index febd7969..bb1efb9b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -51,7 +51,7 @@ dev = [
     "pytest-cov",
     "coverage [toml]",
     # code style
-    "ruff==0.8.*",
+    "ruff==0.9.*",
     "mypy",
     # docs
     "mkdocs",


Test patch:
diff --git a/tests/test_annubes_env.py b/tests/test_annubes_env.py
index 9fb5a95b..2661aef0 100644
--- a/tests/test_annubes_env.py
+++ b/tests/test_annubes_env.py
@@ -68,42 +68,42 @@ def test_fix_time_types(
 
         # Check if the sampled time is in the given list of values
         if isinstance(time, list):
-            assert (
-                env._duration[t] in time
-            ), f"Expected {t} time to be one of {time}, but got {env._duration['fixation']}"
+            assert env._duration[t] in time, (
+                f"Expected {t} time to be one of {time}, but got {env._duration['fixation']}"
+            )
         # Check if the sampled time is in the given range
         elif isinstance(time, tuple) and time[0] == "uniform":
             time_range = time[1]
             assert (
                 time_range[0] <= env._duration[t] <= time_range[1]
             ), f"""Expected {t} time to be between {time_range[0]} and {time_range[1]},
-            but got {env._duration['fixation']}"""
+            but got {env._duration["fixation"]}"""
         # Check if the sampled time is in the given list of values
         elif isinstance(time, tuple) and time[0] == "choice":
-            assert (
-                env._duration[t] in time[1]
-            ), f"Expected {t} time to be one of {time[1]}, but got {env._duration['fixation']}"
+            assert env._duration[t] in time[1], (
+                f"Expected {t} time to be one of {time[1]}, but got {env._duration['fixation']}"
+            )
         # Check if the sampled time is in the given range
         elif isinstance(time, tuple) and time[0] == "truncated_exponential":
             time_range = time[1]
             assert (
                 time_range[0] <= env._duration[t] <= time_range[1]
             ), f"""Expected {t} time to be between {time_range[0]} and {time_range[1]},
-            but got {env._duration['fixation']}"""
+            but got {env._duration["fixation"]}"""
         # Check if the sampled time is the given constant value
         elif isinstance(time, tuple) and time[0] == "constant":
-            assert (
-                env._duration[t] == time[1]
-            ), f"Expected {t} time to be {time[1]}, but got {env._duration['fixation']}"
+            assert env._duration[t] == time[1], (
+                f"Expected {t} time to be {time[1]}, but got {env._duration['fixation']}"
+            )
 
         # For callable time, check if it's actually called
         if callable(time):
             assert env._duration[t] == 500, f"Expected {t} time to be 500, but got {env._duration['fixation']}"
 
         # Ensure the sampled time is a multiple of dt
-        assert (
-            env._duration[t] % env.dt == 0
-        ), f"Expected {t} time to be a multiple of dt ({env.dt}), but got {env._duration['fixation']}"
+        assert env._duration[t] % env.dt == 0, (
+            f"Expected {t} time to be a multiple of dt ({env.dt}), but got {env._duration['fixation']}"
+        )
 
 
 @pytest.mark.parametrize("catch_prob", [0.0, 0.3, 0.7, 1.0])
@@ -179,16 +179,16 @@ def test_annubes_env_max_sequential(session: dict, catch_prob: float, max_sequen
         count = trial_types.count(trial_type)
         expected_count = expected_non_catch * expected_prob
         # Allow for 20% deviation from expected count
-        assert (
-            0.7 * expected_count <= count <= 1.3 * expected_count
-        ), f"{trial_type} trials are not balanced. Expected: {expected_count}, Actual: {count}"
+        assert 0.7 * expected_count <= count <= 1.3 * expected_count, (
+            f"{trial_type} trials are not balanced. Expected: {expected_count}, Actual: {count}"
+        )
 
     # Check catch trial frequency
     catch_count = trial_types.count(None)
     expected_catch = N_TRIALS * catch_prob
-    assert (
-        0.7 * expected_catch <= catch_count <= 1.3 * expected_catch
-    ), f"Catch trials are not balanced. Expected: {expected_catch}, Actual: {catch_count}"
+    assert 0.7 * expected_catch <= catch_count <= 1.3 * expected_catch, (
+        f"Catch trials are not balanced. Expected: {expected_catch}, Actual: {catch_count}"
+    )
 
 
 def test_observation_space(default_env: AnnubesEnv, custom_env: AnnubesEnv) -> None:
@@ -284,9 +284,9 @@ def test_initial_state_and_first_reset(default_env: AnnubesEnv) -> None:
     assert default_env.t == 100, f"default_env={default_env.t}, should be 100 after first reset"
     assert default_env.num_tr == 1, f"num_tr={default_env.num_tr}, should be 1 after first reset"
     assert isinstance(default_env.trial, dict)
-    assert (
-        ob.shape == default_env.observation_space.shape
-    ), f"observation_space.shape={default_env.observation_space.shape}, should match the observation space"
+    assert ob.shape == default_env.observation_space.shape, (
+        f"observation_space.shape={default_env.observation_space.shape}, should match the observation space"
+    )
 
 
 def test_random_seed_reproducibility() -> None:
diff --git a/tests/test_random.py b/tests/test_random.py
index 7b1fd4f8..96ca2e80 100644
--- a/tests/test_random.py
+++ b/tests/test_random.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from neurogym.utils.random import TruncExp
+from neurogym.utils.ngym_random import TruncExp
 from neurogym.utils.scheduler import RandomBlockSchedule, RandomSchedule
 
 
diff --git a/tests/test_wrappers.py b/tests/test_wrappers.py
index dcd49e5b..5b26af6d 100644
--- a/tests/test_wrappers.py
+++ b/tests/test_wrappers.py
@@ -54,11 +54,6 @@ def test_sidebias(
         For example: [(.005, .005, .99), (.005, .99, .005), (.99, .005, .005)],
         corresponds to 3 blocks with each of them giving 0.99 probabilitiy to
         ground truth 3, 2 and 1, respectively.
-
-    Returns:
-    -------
-    None.
-
     """
     if probs is None:
         probs = [(0.005, 0.005, 0.99), (0.005, 0.99, 0.005), (0.99, 0.005, 0.005)]
@@ -102,11 +97,6 @@ def test_passaction(
         number of steps to run the environment (1000)
     verbose : boolean, optional
         whether to print observation and action (False)
-
-    Returns:
-    -------
-    None.
-
     """
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message=".*get variables from other wrappers is deprecated*")
@@ -146,11 +136,6 @@ def test_passreward(
         number of steps to run the environment (1000)
     verbose : boolean, optional
         whether to print observation and reward (False)
-
-    Returns:
-    -------
-    None.
-
     """
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", message=".*get variables from other wrappers is deprecated*")
@@ -197,11 +182,6 @@ def test_reactiontime(
         whether to print observation and reward (False)
     ths : list, optional
         list containing the threholds to make a decision ([-.5, .5])
-
-    Returns:
-    -------
-    None.
-
     """
     if ths is None:
         ths = [-0.5, 0.5]
@@ -299,11 +279,6 @@ def test_variablemapping(
         number of channels (4)
     margin : float, optional
         margin allowed when comparing actual and expected mean block durations (2)
-
-    Returns:
-    -------
-    None.
-
     """
     env_args = {
         "n_stims": n_stims,
@@ -375,9 +350,9 @@ def test_variablemapping(
     assert (sess_durs > min_sess_dur).all()
     mean_sess_dur = np.mean(sess_durs)
     exp_sess_durs = min_sess_dur + 1 / sess_end_prob
-    assert (
-        np.abs(mean_sess_dur - exp_sess_durs) < margin
-    ), f"Mean sess. dur.: {mean_sess_dur}, expected: {1 / sess_end_prob}"
+    assert np.abs(mean_sess_dur - exp_sess_durs) < margin, (
+        f"Mean sess. dur.: {mean_sess_dur}, expected: {1 / sess_end_prob}"
+    )
     mapp_blck_durs = np.diff(mapp_ch)
     assert (mapp_blck_durs > min_mapp_dur).all()
     mean_durs = np.mean(mapp_blck_durs)
@@ -411,11 +386,6 @@ def test_noise(
         margin allowed when comparing actual and expected performances (0.01)
     perf_th : float, optional
         target performance for the noise wrapper (0.7)
-
-    Returns:
-    -------
-    None.
-
     """
     env_args = {"timing": {"fixation": 100, "stimulus": 200, "decision": 200}}
     env = gym.make(env, **env_args)