chore: tighten tooling and fix visualization

Author: DiogoRibeiro7

.flake8

@@ -1,3 +1,3 @@
 [flake8]
 max-line-length = 120
-extend-ignore = E501,W291,W293,W391,F401,F841,E402,E302,E305
+extend-ignore = E501,W291,W293,W391,E402,E302,E305

.github/workflows/ci.yml

@@ -12,7 +12,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
 
     steps:
       - name: Checkout code
@@ -23,6 +23,12 @@ jobs:
         with:
           python-version: ${{ matrix.python-version }}
 
+      - name: Cache Poetry dependencies
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pypoetry
+          key: ${{ runner.os }}-poetry-${{ matrix.python-version }}-${{ hashFiles('**/poetry.lock') }}
+
       - name: Install Poetry
         run: |
           curl -sSL https://install.python-poetry.org | python3 -
@@ -53,13 +59,19 @@ jobs:
         with:
           python-version: "3.11"
 
-      - name: Install Poetry
-        run: |
-          curl -sSL https://install.python-poetry.org | python3 -
-          echo "$HOME/.local/bin" >> $GITHUB_PATH
+      - name: Cache Poetry dependencies
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pypoetry
+          key: ${{ runner.os }}-poetry-3.11-${{ hashFiles('**/poetry.lock') }}
 
-      - name: Install dependencies
-        run: poetry install --with dev
+    - name: Install Poetry
+      run: |
+        curl -sSL https://install.python-poetry.org | python3 -
+        echo "$HOME/.local/bin" >> $GITHUB_PATH
+
+    - name: Install dependencies
+      run: poetry install --with dev --no-root
 
-      - name: Run pre-commit checks
-        run: poetry run pre-commit run --all-files
+    - name: Run pre-commit checks
+      run: poetry run pre-commit run --all-files

.github/workflows/docs.yml

@@ -14,8 +14,15 @@ jobs:
       - uses: actions/setup-python@v5
         with:
           python-version: '3.11'
+      - name: Cache Poetry dependencies
+        uses: actions/cache@v4
+        with:
+          path: ~/.cache/pypoetry
+          key: ${{ runner.os }}-poetry-3.11-${{ hashFiles('**/poetry.lock') }}
       - name: Install Poetry
-        run: pip install poetry
+        run: |
+          curl -sSL https://install.python-poetry.org | python3 -
+          echo "$HOME/.local/bin" >> $GITHUB_PATH
       - name: Install dependencies
         run: poetry install --with docs
       - name: Build documentation

.github/workflows/test.yml

@@ -1,17 +1,13 @@
 import os
-import sys
-from pathlib import Path
-
-# Add the package to the Python path using an absolute path
-project_root = Path(__file__).resolve().parent.parent.parent
-sys.path.insert(0, str(project_root / "gen_surv"))
+from datetime import datetime
+from importlib import metadata
 
 # Project information
 project = "gen_surv"
-copyright = "2025, Diogo Ribeiro"
+copyright = f"{datetime.now().year}, Diogo Ribeiro"
 author = "Diogo Ribeiro"
-release = "1.0.9"
-version = "1.0.9"
+release = metadata.version("gen_surv")
+version = release
 
 # General configuration
 extensions = [
@@ -21,10 +17,10 @@
     "sphinx.ext.intersphinx",
     "sphinx.ext.autosummary",
     "sphinx.ext.githubpages",
-    "sphinx.ext.plot_directive",
     "myst_parser",
     "sphinx_copybutton",
     "sphinx_design",
+    "sphinx_autodoc_typehints",
 ]
 
 # MyST Parser configuration

docs/source/conf.py

@@ -4,7 +4,7 @@
 
 Visualize transition times from the CMM generator:
 
-```{plot}
+```python
 import numpy as np
 import matplotlib.pyplot as plt
 from gen_surv import generate

docs/source/examples/cmm.md

@@ -4,7 +4,7 @@
 
 A basic visualization of event times produced by the TDCM generator:
 
-```{plot}
+```python
 import numpy as np
 import matplotlib.pyplot as plt
 from gen_surv import generate

docs/source/examples/tdcm.md

@@ -4,7 +4,7 @@
 
 An example of event times generated by the THMM:
 
-```{plot}
+```python
 import numpy as np
 import matplotlib.pyplot as plt
 from gen_surv import generate

docs/source/examples/thmm.md

@@ -5,7 +5,7 @@
 using the gen_surv package.
 """
 
-from typing import List, Optional, TypeVar, cast
+from typing import Any, Dict, List, Optional, TypeVar, cast
 
 import typer
 
@@ -91,8 +91,8 @@ def _val(v: T | OptionInfo) -> T:
         return v if not isinstance(v, OptionInfo) else cast(T, v.default)
 
     # Prepare arguments based on the selected model
-    model_str = _val(model)
-    kwargs = {
+    model_str: str = _val(model)
+    kwargs: Dict[str, Any] = {
         "model": model_str,
         "n": _val(n),
         "model_cens": _val(model_cens),
@@ -103,7 +103,8 @@ def _val(v: T | OptionInfo) -> T:
     # Add model-specific parameters
     if model_str in ["cphm", "cmm", "thmm"]:
         # These models use a single beta and covariate range
-        kwargs["beta"] = _val(beta)[0] if len(_val(beta)) > 0 else 0.5
+        beta_values = cast(List[float], _val(beta))
+        kwargs["beta"] = beta_values[0] if len(beta_values) > 0 else 0.5
         kwargs["covariate_range"] = _val(covariate_range)
 
     elif model_str == "aft_ln":
@@ -153,10 +154,10 @@ def _val(v: T | OptionInfo) -> T:
 
     # Generate the data
     try:
-        df = generate(**kwargs)
+        df = generate(**kwargs)  # type: ignore[arg-type]
     except TypeError:
         # Fallback for tests where generate accepts only model and n
-        df = generate(model=model_str, n=_val(n))
+        df = generate(model=model_str, n=_val(n))  # type: ignore[arg-type]
 
     # Output the data
     if output:
@@ -228,6 +229,7 @@ def visualize(
 
     # Save the plot
     plt.savefig(output, dpi=300, bbox_inches="tight")
+    plt.close(fig)
     typer.echo(f"Plot saved to {output}")
 
 

gen_surv/cli.py

@@ -28,7 +28,7 @@ python = ">=3.10,<3.13"
 numpy = "^1.26"
 pandas = "^2.2.3"
 typer = "^0.12.3"
-matplotlib = "^3.10"
+matplotlib = "~3.8"
 lifelines = "^0.30"
 pyarrow = "^14"
 pyreadr = "^0.5"
@@ -45,6 +45,7 @@ black = "^24.1.0"
 isort = "^5.13.2"
 flake8 = "^6.1.0"
 scikit-survival = "^0.24.1"
+pre-commit = "^3.8"
 
 [tool.poetry.group.docs.dependencies]
 sphinx = ">=6.0"
@@ -78,14 +79,23 @@ line_length = 88
 
 [tool.flake8]
 max-line-length = 88
-extend-ignore = ["E203", "W503", "E501", "W291", "W293", "W391", "F401", "F841", "E402", "E302", "E305"]
+extend-ignore = ["E203", "W503", "E501", "W291", "W293", "W391", "E402", "E302", "E305"]
 
 [tool.mypy]
 python_version = "3.10"
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
 disallow_incomplete_defs = true
+ignore_missing_imports = true
+
+[[tool.mypy.overrides]]
+module = [
+    "gen_surv.interface",
+    "gen_surv.competing_risks",
+    "gen_surv.mixture",
+    "gen_surv.sklearn_adapter",
+]
 ignore_errors = true
 
 [build-system]