fix: fix black and mypi issues

Author: DiogoRibeiro7

gen_surv/__init__.py

@@ -6,8 +6,6 @@
 from importlib.metadata import PackageNotFoundError, version
 
 from .aft import gen_aft_log_logistic, gen_aft_log_normal, gen_aft_weibull
-
-# Helper functions
 from .bivariate import sample_bivariate_distribution
 from .censoring import (
     CensoringModel,
@@ -22,54 +20,43 @@
 )
 from .cmm import gen_cmm
 from .competing_risks import gen_competing_risks, gen_competing_risks_weibull
-
-# Individual generators
 from .cphm import gen_cphm
 from .export import export_dataset
-from .integration import to_sksurv
-
-# Main interface
 from .interface import generate
 from .mixture import cure_fraction_estimate, gen_mixture_cure
 from .piecewise import gen_piecewise_exponential
 from .sklearn_adapter import GenSurvDataGenerator
 from .tdcm import gen_tdcm
 from .thmm import gen_thmm
 
+# Get package version
+try:
+    __version__ = version("gen_surv")
+except PackageNotFoundError:  # pragma: no cover - fallback when package not installed
+    __version__ = "0.0.0"
+
 # Visualization tools (requires matplotlib and lifelines)
 try:
-    from .visualization import describe_survival  # noqa: F401
-    from .visualization import plot_covariate_effect  # noqa: F401
-    from .visualization import plot_hazard_comparison  # noqa: F401
-    from .visualization import plot_survival_curve  # noqa: F401
+    from .visualization import (  # noqa: F401
+        describe_survival,
+        plot_covariate_effect,
+        plot_hazard_comparison,
+        plot_survival_curve,
+    )
 
     _has_visualization = True
 except ImportError:
     _has_visualization = False
 
-"""Top-level package for ``gen_surv``.
-
-This module exposes the :func:`generate` function and provides access to the
-package version via ``__version__``.
-"""
-
-from importlib.metadata import PackageNotFoundError, version
-
-from .interface import generate
-
+# Optional scikit-survival integration
 try:
-    __version__ = version("gen_surv")
-except PackageNotFoundError:  # pragma: no cover - fallback when package not installed
-    __version__ = "0.0.0"
+    from .integration import from_sksurv, to_sksurv  # noqa: F401
 
-# Optional imports - only available if dependencies are installed
-try:
-    from .integration import to_sksurv, from_sksurv
-    __all__ = ["generate", "__version__", "to_sksurv", "from_sksurv"]
+    _has_sksurv = True
 except ImportError:
-    # scikit-survival not available
-    __all__ = ["generate", "__version__"]
+    _has_sksurv = False
 
+# Define exports
 __all__ = [
     # Main interface
     "generate",
@@ -87,7 +74,7 @@
     "gen_mixture_cure",
     "cure_fraction_estimate",
     "gen_piecewise_exponential",
-    # Helpers
+    # Helper functions
     "sample_bivariate_distribution",
     "runifcens",
     "rexpocens",
@@ -99,11 +86,13 @@
     "GammaCensoring",
     "CensoringModel",
     "export_dataset",
-    "to_sksurv",
     "GenSurvDataGenerator",
 ]
 
-# Add visualization tools to __all__ if available
+# Add optional exports if available
+if _has_sksurv:
+    __all__.extend(["to_sksurv", "from_sksurv"])
+
 if _has_visualization:
     __all__.extend(
         [

gen_surv/integration.py

@@ -5,15 +5,18 @@
 
 try:
     from sksurv.util import Surv
+
     SKSURV_AVAILABLE = True
 except ImportError:
     SKSURV_AVAILABLE = False
 
 
-def to_sksurv(df: pd.DataFrame, time_col="time", event_col="status"):
+def to_sksurv(
+    df: pd.DataFrame, time_col: str = "time", event_col: str = "status"
+) -> np.ndarray:
     """
     Convert a pandas DataFrame to a scikit-survival structured array.
-    
+
     Parameters
     ----------
     df : pd.DataFrame
@@ -22,12 +25,12 @@ def to_sksurv(df: pd.DataFrame, time_col="time", event_col="status"):
         Name of the column containing survival times
     event_col : str, default "status"
         Name of the column containing event indicators (0/1 or boolean)
-        
+
     Returns
     -------
     y : structured array
         Structured array suitable for scikit-survival functions
-        
+
     Raises
     ------
     ImportError
@@ -37,49 +40,50 @@ def to_sksurv(df: pd.DataFrame, time_col="time", event_col="status"):
     """
     if not SKSURV_AVAILABLE:
         raise ImportError("scikit-survival is required but not installed")
-    
+
     if df.empty:
         # Handle empty DataFrame case by creating a minimal valid structured array
         # This avoids the "event indicator must be binary" error for empty arrays
         return np.array([], dtype=[(event_col, bool), (time_col, float)])
-    
+
     if time_col not in df.columns:
         raise ValueError(f"Column '{time_col}' not found in DataFrame")
     if event_col not in df.columns:
         raise ValueError(f"Column '{event_col}' not found in DataFrame")
-    
+
     return Surv.from_dataframe(event_col, time_col, df)
 
 
-def from_sksurv(y: np.ndarray, time_col="time", event_col="status"):
+def from_sksurv(
+    y: np.ndarray, time_col: str = "time", event_col: str = "status"
+) -> pd.DataFrame:
     """
     Convert a scikit-survival structured array to a pandas DataFrame.
-    
+
     Parameters
     ----------
     y : structured array
         Structured array from scikit-survival
     time_col : str, default "time"
         Name for the time column in the resulting DataFrame
-    event_col : str, default "status"  
+    event_col : str, default "status"
         Name for the event column in the resulting DataFrame
-        
+
     Returns
     -------
     df : pd.DataFrame
         DataFrame with time and event columns
     """
     if not SKSURV_AVAILABLE:
         raise ImportError("scikit-survival is required but not installed")
-    
+
     if len(y) == 0:
         return pd.DataFrame({time_col: [], event_col: []})
-    
+
     # Extract field names from structured array
     event_field = y.dtype.names[0]
     time_field = y.dtype.names[1]
-    
-    return pd.DataFrame({
-        time_col: y[time_field],
-        event_col: y[event_field].astype(int)
-    })
+
+    return pd.DataFrame(
+        {time_col: y[time_field], event_col: y[event_field].astype(int)}
+    )

tests/test_integration_sksurv.py

@@ -4,20 +4,17 @@
 import pandas as pd
 import pytest
 
-from gen_surv.integration import to_sksurv, from_sksurv
+from gen_surv.integration import from_sksurv, to_sksurv
 
 
 def test_to_sksurv_basic():
     """Test basic conversion from DataFrame to sksurv format."""
     pytest.importorskip("sksurv.util")
-    
-    df = pd.DataFrame({
-        "time": [1.0, 2.0, 3.0],
-        "status": [1, 0, 1]
-    })
-    
+
+    df = pd.DataFrame({"time": [1.0, 2.0, 3.0], "status": [1, 0, 1]})
+
     arr = to_sksurv(df)
-    
+
     assert len(arr) == 3
     assert arr.dtype.names == ("status", "time")
     assert list(arr["time"]) == [1.0, 2.0, 3.0]
@@ -27,49 +24,48 @@ def test_to_sksurv_basic():
 def test_to_sksurv_custom_columns():
     """Test conversion with custom column names."""
     pytest.importorskip("sksurv.util")
-    
-    df = pd.DataFrame({
-        "survival_time": [1.0, 2.0],
-        "event": [1, 0]
-    })
-    
+
+    df = pd.DataFrame({"survival_time": [1.0, 2.0], "event": [1, 0]})
+
     arr = to_sksurv(df, time_col="survival_time", event_col="event")
-    
+
     assert len(arr) == 2
     assert arr.dtype.names == ("event", "survival_time")
 
 
 def test_to_sksurv_empty_dataframe():
     """Test conversion of empty DataFrame."""
     pytest.importorskip("sksurv.util")
-    
+
     df = pd.DataFrame({"time": [], "status": []})
     arr = to_sksurv(df)
-    
+
     assert len(arr) == 0
     assert arr.dtype.names == ("status", "time")
 
 
 def test_to_sksurv_missing_columns():
     """Test error handling for missing columns."""
     pytest.importorskip("sksurv.util")
-    
+
     df = pd.DataFrame({"time": [1.0, 2.0]})
-    
+
     with pytest.raises(ValueError, match="Column 'status' not found"):
         to_sksurv(df)
 
 
 def test_from_sksurv_basic():
     """Test conversion from sksurv format to DataFrame."""
     pytest.importorskip("sksurv.util")
-    
+
     # Create a structured array manually
-    arr = np.array([(True, 1.0), (False, 2.0), (True, 3.0)], 
-                   dtype=[("status", bool), ("time", float)])
-    
+    arr = np.array(
+        [(True, 1.0), (False, 2.0), (True, 3.0)],
+        dtype=[("status", bool), ("time", float)],
+    )
+
     df = from_sksurv(arr)
-    
+
     assert len(df) == 3
     assert list(df.columns) == ["time", "status"]
     assert list(df["time"]) == [1.0, 2.0, 3.0]
@@ -79,27 +75,24 @@ def test_from_sksurv_basic():
 def test_from_sksurv_empty():
     """Test conversion of empty structured array."""
     pytest.importorskip("sksurv.util")
-    
+
     arr = np.array([], dtype=[("status", bool), ("time", float)])
     df = from_sksurv(arr)
-    
+
     assert len(df) == 0
     assert list(df.columns) == ["time", "status"]
 
 
 def test_roundtrip_conversion():
     """Test that conversion is bidirectional."""
     pytest.importorskip("sksurv.util")
-    
-    original_df = pd.DataFrame({
-        "time": [1.0, 2.5, 4.0],
-        "status": [1, 0, 1]
-    })
-    
+
+    original_df = pd.DataFrame({"time": [1.0, 2.5, 4.0], "status": [1, 0, 1]})
+
     # Convert to sksurv and back
     arr = to_sksurv(original_df)
     result_df = from_sksurv(arr)
-    
+
     pd.testing.assert_frame_equal(original_df, result_df)