Feat/consolidate and optimize (#63)

* chore: tighten tooling and fix visualization

* docs: document Python 3.10 requirement

* ci: align workflow with review template

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

@@ -2,17 +2,19 @@ name: CI
 
 on:
   push:
-    branches: [main]
+    branches:
+      - main
   pull_request:
-    branches: [main]
+    branches:
+      - main
 
 jobs:
   test:
     name: Test with Python ${{ matrix.python-version }}
     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 +25,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 -
@@ -38,7 +46,7 @@ jobs:
         uses: codecov/codecov-action@v5
         with:
           files: coverage.xml
-          token: ${{ secrets.CODECOV_TOKEN }}  # optional if public repo
+          token: ${{ secrets.CODECOV_TOKEN }}
 
   lint:
     name: Code Quality
@@ -53,13 +61,19 @@ jobs:
         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: |
           curl -sSL https://install.python-poetry.org | python3 -
           echo "$HOME/.local/bin" >> $GITHUB_PATH
 
       - name: Install dependencies
-        run: poetry install --with dev
+        run: poetry install --with dev --no-root
 
       - 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

@@ -3,7 +3,7 @@
 [![Coverage](https://codecov.io/gh/DiogoRibeiro7/genSurvPy/branch/main/graph/badge.svg)](https://app.codecov.io/gh/DiogoRibeiro7/genSurvPy)
 [![Docs](https://readthedocs.org/projects/gensurvpy/badge/?version=latest)](https://gensurvpy.readthedocs.io/en/latest/)
 [![PyPI](https://img.shields.io/pypi/v/gen_surv)](https://pypi.org/project/gen-surv/)
-[![Tests](https://github.com/DiogoRibeiro7/genSurvPy/actions/workflows/test.yml/badge.svg)](https://github.com/DiogoRibeiro7/genSurvPy/actions/workflows/test.yml)
+[![Tests](https://github.com/DiogoRibeiro7/genSurvPy/actions/workflows/ci.yml/badge.svg)](https://github.com/DiogoRibeiro7/genSurvPy/actions/workflows/ci.yml)
 [![Python](https://img.shields.io/pypi/pyversions/gen_surv)](https://pypi.org/project/gen-surv/)
 
 **gen_surv** is a Python package for simulating survival data under a variety of statistical models. It is inspired by the R package [genSurv](https://cran.r-project.org/package=genSurv) and provides a unified interface for generating realistic survival datasets.
@@ -25,6 +25,8 @@
 
 ## Installation
 
+Requires Python 3.10 or later.
+
 Install the latest release from PyPI:
 
 ```bash

README.md

@@ -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}")