Fix mypy and pre-commit hooks

.pre-commit-config.yaml

@@ -9,16 +9,6 @@ repos:
     rev: 1.16.0
     hooks:
       - id: yamlfix
-  - repo: local
-    hooks:
-      - id: check-nbqa-version-mismatch
-        name: Check for version mismatch between black, ruff, and nbQA
-        entry: python scripts/check_nbqa_version_mismatch.py
-        language: python
-        always_run: true
-        require_serial: true
-        additional_dependencies:
-          - pyyaml
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.6.0
     hooks:
@@ -55,30 +45,27 @@ repos:
     rev: v1.35.1
     hooks:
       - id: yamllint
-  - repo: https://github.com/psf/black
-    rev: 24.4.2
-    hooks:
-      - id: black
-        language_version: python3.12
   - repo: https://github.com/asottile/blacken-docs
     rev: 1.18.0
     hooks:
       - id: blacken-docs
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.5.0
+    rev: v0.5.1
     hooks:
+      # Run the linter.
       - id: ruff
-      #   args:
-      #     - --verbose
-  - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.8.5
-    hooks:
-      - id: nbqa-black
-        additional_dependencies:
-          - black==24.4.2
-      - id: nbqa-ruff
-        additional_dependencies:
-          - ruff==v0.5.0
+        types_or:
+          - python
+          - pyi
+          - jupyter
+        args:
+          - --fix
+      # Run the formatter.
+      - id: ruff-format
+        types_or:
+          - python
+          - pyi
+          - jupyter
   - repo: https://github.com/executablebooks/mdformat
     rev: 0.7.17
     hooks:

explanations/dispatchers.ipynb

@@ -682,10 +682,8 @@
     "        strict=False,\n",
     "    ),\n",
     "):\n",
-    "\n",
     "    # loop over product of dense variables\n",
     "    for j, wealth in enumerate(sc_space.dense_vars[\"wealth\"]):\n",
-    "\n",
     "        u = utility(\n",
     "            wealth=wealth,\n",
     "            retirement=retirement,\n",

pyproject.toml

@@ -164,6 +164,7 @@ test-gpu = {features = ["test", "cuda"], solve-group = "cuda"}
 [tool.ruff]
 target-version = "py312"
 fix = true
+exclude = ["src/lcm/sandbox"]
 
 [tool.ruff.lint]
 select = ["ALL"]
@@ -215,6 +216,9 @@ extend-ignore = [
 
     # long messages outside the exception class
     "TRY003",
+
+    # Missing docstring in magic method
+    "D105",
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -223,6 +227,7 @@ extend-ignore = [
 "examples/*" = ["INP001"]
 "explanations/*" = ["INP001", "B018", "T201", "E402", "PD008"]
 "scripts/*" = ["INP001", "D101", "RET503"]
+"**/*.ipynb" = ["FBT003", "E402"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "google"
@@ -238,7 +243,6 @@ black = "pyproject.toml"
 black = 1
 
 [tool.nbqa.exclude]
-ruff = "src/lcm/sandbox"
 black = "src/lcm/sandbox"
 
 # ======================================================================================

src/lcm/__init__.py

@@ -1,5 +1,4 @@
 from lcm import mark
-from lcm.user_grids import DiscreteGrid, LinspaceGrid, LogspaceGrid
-from lcm.user_model import Model
+from lcm.user_input import DiscreteGrid, LinspaceGrid, LogspaceGrid, Model
 
 __all__ = ["mark", "Model", "LinspaceGrid", "LogspaceGrid", "DiscreteGrid"]

src/lcm/create_params_template.py

@@ -7,7 +7,7 @@
 from jax import Array
 
 from lcm.typing import Params, ScalarUserInput
-from lcm.user_model import Model
+from lcm.user_input import Model
 
 
 def create_params_template(
@@ -129,7 +129,6 @@ def _create_stochastic_transition_params(
     invalid_dependencies = {}
 
     for var in stochastic_variables:
-
         # Retrieve corresponding next function and its arguments
         next_var = user_model.functions[f"next_{var}"]
         dependencies = list(inspect.signature(next_var).parameters)

src/lcm/entry_point.py

@@ -18,7 +18,7 @@
 from lcm.simulate import simulate
 from lcm.solve_brute import solve
 from lcm.state_space import create_state_choice_space
-from lcm.user_model import Model
+from lcm.user_input import Model
 
 
 def get_lcm_function(

src/lcm/process_model.py

@@ -15,13 +15,19 @@
 from lcm.interfaces import (
     ContinuousGridInfo,
     ContinuousGridSpec,
+    ContinuousGridType,
     DiscreteGridSpec,
     GridSpec,
     InternalModel,
 )
 from lcm.typing import Params
-from lcm.user_grids import ContinuousGrid, DiscreteGrid, LinspaceGrid, LogspaceGrid
-from lcm.user_model import Model
+from lcm.user_input import (
+    ContinuousGrid,
+    DiscreteGrid,
+    LinspaceGrid,
+    LogspaceGrid,
+    Model,
+)
 
 
 def process_model(user_model: Model) -> InternalModel:
@@ -225,7 +231,7 @@ def _get_gridspecs(
     for name, spec in raw_variables.items():
         if isinstance(spec, ContinuousGrid):
             if isinstance(spec, LinspaceGrid):
-                kind = "linspace"
+                kind: ContinuousGridType = "linspace"
             elif isinstance(spec, LogspaceGrid):
                 kind = "logspace"
             else: