102 separate linting workflow in ci

.github/workflows/linting_formatting.yml

@@ -0,0 +1,27 @@
+name: Linting and Formatting
+
+on:
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v7
+        with:
+          persist-credentials: false
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          version: "0.12.0"
+
+      - name: Install package
+        run: uv sync --group dev
+
+      - name: Check styling
+        run: uv run ruff format --check .
+
+      - name: Check linting
+        run: uv run ruff check src/lanfactory

.github/workflows/run_tests.yml

@@ -19,24 +19,20 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: ${{ matrix.python-version }}
-
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:
-          version: "0.6.5"
+          version: "0.12.0"
+          python-version: ${{ matrix.python-version }}
           # enable-cache: true
           # cache-dependency-glob: "pyproject.toml pdm.lock"
 
       - name: Clear all caches
         run: |
-            rm -rf ~/.cache/pip
-            rm -rf ~/.cache/uv
-            rm -rf ~/.cache/conda
-            rm -rf ~/.cache/npm
+          rm -rf ~/.cache/pip
+          rm -rf ~/.cache/uv
+          rm -rf ~/.cache/conda
+          rm -rf ~/.cache/npm
 
       - name: Install package
         run: uv sync --all-groups --reinstall
@@ -47,12 +43,6 @@ jobs:
       - name: Run pytest
         run: uv run pytest
 
-      - name: Check styling
-        run: uv run ruff format --check .
-
-      - name: Linting
-        run: uv run ruff check src/lanfactory
-
       - name: Upload coverage reports to Codecov
         uses: codecov/codecov-action@v7
         with:
@@ -67,15 +57,11 @@ jobs:
         with:
           persist-credentials: false
 
-      - name: Setup Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: "3.12"
-
       - name: Install uv
         uses: astral-sh/setup-uv@v7
         with:
-          version: "0.6.5"
+          version: "0.12.0"
+          python-version: ${{ matrix.python-version }}
 
       - name: Install package (with notebook + backend deps)
         run: uv sync --all-groups

.gitignore

@@ -157,3 +157,4 @@ explorations/
 __marimo__/
 
 uv.lock
+.vscode/

pyproject.toml

@@ -82,7 +82,7 @@ dev = [
     "pytest-timer>=1.0.0",
     "pytest-xdist>=3.6.1",
     "pytest>=8.3.1",
-    "ruff>=0.14.4",
+    "ruff>=0.15.1",
     "types-PyYAML",
     "mlflow>=3.14.0",
     "jaxonnxruntime>=0.3",
@@ -120,6 +120,14 @@ line-length = 88
 src = ["src/lanfactory", "tests"]
 exclude = ["notebooks/*", "docs/*", "docs/basic_tutorial/basic_tutorial.ipynb"]
 
+[tool.ruff.lint]
+ignore = [
+    "B008",   # Function calls in default arguments (Typer/FastAPI options)
+    "EXE001", # Shebang present but file is not executable
+    "EXE002", # Executable file fails to declare a shebang
+    "BLE001", #
+]
+
 # Coverage settings
 [tool.coverage.run]
 source = ["src"]

src/lanfactory/__init__.py

@@ -1,11 +1,8 @@
 __version__ = "0.8.0"
 
-from . import config
-from . import trainers
-from . import utils
-from . import onnx
+from . import config, onnx, trainers, utils
 
-__all__ = ["config", "trainers", "utils", "onnx", "network_inspectors"]
+__all__ = ["config", "network_inspectors", "onnx", "trainers", "utils"]
 
 
 def __getattr__(name):

src/lanfactory/cli/download_hf.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 
 import typer
-
 from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES
 
 app = typer.Typer()

src/lanfactory/cli/jax_train.py

@@ -164,9 +164,10 @@ def main(
 
     if mlflow_tracking_enabled:
         try:
-            import mlflow
             import os
 
+            import mlflow
+
             # Set tracking URI with priority: CLI arg > env var > default
             if mlflow_tracking_uri:
                 tracking_uri = mlflow_tracking_uri
@@ -234,9 +235,10 @@ def main(
 
             if not mlflow_tracking_enabled:
                 # Need to initialize MLflow just for querying
-                import mlflow
                 import os
 
+                import mlflow
+
                 # Use same logic as above for tracking URI
                 if mlflow_tracking_uri:
                     tracking_uri = mlflow_tracking_uri
@@ -312,7 +314,7 @@ def main(
     # Mode 2: Validation - verify MLflow files exist in training_data_folder
     if mlflow_lineage_info and training_data_folder:
         expected_files = set(mlflow_lineage_info["all_files"])
-        actual_files = set(f.name for f in valid_file_list)
+        actual_files = {f.name for f in valid_file_list}
 
         missing_files = expected_files - actual_files
         extra_files = actual_files - expected_files
@@ -463,13 +465,8 @@ def main(
         ]
     )
 
-    pickle.dump(
-        network_config,
-        open(
-            networks_path / file_name_suffix,
-            "wb",
-        ),
-    )
+    file_path = networks_path / file_name_suffix
+    file_path.write_bytes(pickle.dumps(network_config))
 
     # Load network
     net = lanfactory.trainers.JaxMLPFactory(

src/lanfactory/cli/torch_train.py

@@ -22,11 +22,10 @@
 from importlib.resources import as_file, files
 from pathlib import Path
 
+import lanfactory
 import psutil
 import torch
 import typer
-
-import lanfactory
 from lanfactory.cli.utils import (
     _get_train_network_config,
 )
@@ -168,9 +167,10 @@ def main(
 
     if mlflow_tracking_enabled:
         try:
-            import mlflow
             import os
 
+            import mlflow
+
             # Set tracking URI with priority: CLI arg > env var > default
             if mlflow_tracking_uri:
                 tracking_uri = mlflow_tracking_uri
@@ -238,9 +238,10 @@ def main(
 
             if not mlflow_tracking_enabled:
                 # Need to initialize MLflow just for querying
-                import mlflow
                 import os
 
+                import mlflow
+
                 # Use same logic as above for tracking URI
                 if mlflow_tracking_uri:
                     tracking_uri = mlflow_tracking_uri
@@ -316,7 +317,7 @@ def main(
     # Mode 2: Validation - verify MLflow files exist in training_data_folder
     if mlflow_lineage_info and training_data_folder:
         expected_files = set(mlflow_lineage_info["all_files"])
-        actual_files = set(f.name for f in valid_file_list)
+        actual_files = {f.name for f in valid_file_list}
 
         missing_files = expected_files - actual_files
         extra_files = actual_files - expected_files
@@ -467,13 +468,8 @@ def main(
         ]
     )
 
-    pickle.dump(
-        network_config,
-        open(
-            networks_path / file_name_suffix,
-            "wb",
-        ),
-    )
+    file_path = networks_path / file_name_suffix
+    file_path.write_bytes(pickle.dumps(network_config))
 
     # Load network
     net = lanfactory.trainers.TorchMLP(

src/lanfactory/cli/upload_hf.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 
 import typer
-
 from lanfactory.hf import DEFAULT_REPO_ID, VALID_NETWORK_TYPES
 
 app = typer.Typer()

src/lanfactory/cli/utils.py

@@ -1,10 +1,11 @@
 # import argparse
 import logging
-from pathlib import Path
 import pickle
-import yaml
-import numpy as np
+from pathlib import Path
+
 import lanfactory
+import numpy as np
+import yaml
 
 logger = logging.getLogger(__name__)
 
@@ -88,9 +89,9 @@ def _make_train_network_configs(
     if save_name:
         save_folder = Path(save_folder)
         save_folder.mkdir(parents=True, exist_ok=True)  # pragma: no cover
-        save_name = save_folder / save_name
-        pickle.dump(config_dict, open(save_name, "wb"))
-        print(f"Saved to: {save_name}")
+        save_path = save_folder / save_name
+        save_path.write_bytes(pickle.dumps(config_dict))
+        print(f"Saved to: {save_path}")
     else:
         print("No save name provided, config not saved to file.")
 
@@ -99,7 +100,7 @@ def _make_train_network_configs(
 
 def _get_train_network_config(yaml_config_path: str | Path | None = None, net_index=0):
     if yaml_config_path is not None:
-        basic_config = yaml.safe_load(open(yaml_config_path, "rb"))
+        basic_config = yaml.safe_load(Path(yaml_config_path).read_bytes())
         network_type = basic_config["NETWORK_TYPE"]
     else:
         raise ValueError("No YAML config path provided")

src/lanfactory/config/__init__.py

@@ -1,21 +1,21 @@
 from .network_configs import (
-    network_config_mlp,
     network_config_choice_prob,
-    network_config_opn,
     network_config_cpn,
-    train_config_mlp,
+    network_config_mlp,
+    network_config_opn,
     train_config_choice_prob,
-    train_config_opn,
     train_config_cpn,
+    train_config_mlp,
+    train_config_opn,
 )
 
 __all__ = [
-    "network_config_mlp",
     "network_config_choice_prob",
-    "network_config_opn",
     "network_config_cpn",
-    "train_config_mlp",
+    "network_config_mlp",
+    "network_config_opn",
     "train_config_choice_prob",
-    "train_config_opn",
     "train_config_cpn",
+    "train_config_mlp",
+    "train_config_opn",
 ]

src/lanfactory/hf/__init__.py

@@ -7,20 +7,20 @@
 DEFAULT_REPO_ID = "franklab/HSSM"
 VALID_NETWORK_TYPES = ("lan", "cpn", "opn")
 
-from lanfactory.hf.model_card import (  # noqa: E402
-    load_model_card_yaml,
-    generate_readme,
+from lanfactory.hf.download import download_model
+from lanfactory.hf.model_card import (
     ModelCardConfig,
+    generate_readme,
+    load_model_card_yaml,
 )
-from lanfactory.hf.upload import upload_model  # noqa: E402
-from lanfactory.hf.download import download_model  # noqa: E402
+from lanfactory.hf.upload import upload_model
 
 __all__ = [
     "DEFAULT_REPO_ID",
     "VALID_NETWORK_TYPES",
-    "load_model_card_yaml",
-    "generate_readme",
     "ModelCardConfig",
-    "upload_model",
     "download_model",
+    "generate_readme",
+    "load_model_card_yaml",
+    "upload_model",
 ]

src/lanfactory/hf/download.py

@@ -123,9 +123,10 @@ def _download_model_hf(  # pragma: no cover
     model_files = [f for f in all_files if f.startswith(path_prefix)]
 
     if not model_files:
+        top_level_dirs = {f.split("/")[0] for f in all_files if "/" in f}
         raise FileNotFoundError(
             f"No files found at {repo_id}/{path_prefix}. "
-            f"Available paths: {set(f.split('/')[0] for f in all_files if '/' in f)}"
+            f"Available paths: {top_level_dirs}"
         )
 
     if include_patterns:

src/lanfactory/hf/upload.py

@@ -155,7 +155,8 @@ def _upload_to_hf(  # pragma: no cover
 ) -> str:
     """HF-dependent implementation of upload_model."""
     try:
-        from huggingface_hub import HfApi, create_repo as hf_create_repo
+        from huggingface_hub import HfApi
+        from huggingface_hub import create_repo as hf_create_repo
     except ImportError as exc:
         raise ImportError(
             "huggingface_hub is required for HuggingFace uploads. "

src/lanfactory/network_inspectors/__init__.py

@@ -15,10 +15,10 @@
 from .loaders import get_torch_mlp
 
 __all__ = [
+    "GridSpec",
+    "ModelSpec",
+    "PlotConfig",
     "get_torch_mlp",
     "kde_vs_lan_likelihoods",
     "lan_manifold",
-    "ModelSpec",
-    "PlotConfig",
-    "GridSpec",
 ]