optional statsmodels

vincentarelbundock · Dec 23, 2023 · ca3046d · ca3046d
1 parent 589a569
commit ca3046d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 deletions.
diff --git a/marginaleffects/sanitize_model.py b/marginaleffects/sanitize_model.py
@@ -2,10 +2,20 @@
 
 
 def sanitize_model(model):
-    # TODO: other than statsmodels
     if model is None:
         return model
 
-    if not isinstance(model, ModelAbstract):
-        model = ModelStatsmodels(model)
-    return model
+    if isinstance(model, ModelAbstract):
+        return model
+
+    try:
+        import statsmodels.base.wrapper as smw
+
+        if isinstance(model, smw.ResultsWrapper):
+            return ModelStatsmodels(model)
+    except ImportError:
+        pass
+
+    raise ValueError(
+        "Unknown model type. Try installing the 'statsmodels' package or file an issue at https://github.com/vincentarelbundock/pymarginaleffects."
+    )
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -8,7 +8,6 @@ readme = "README.md"
 [tool.poetry.dependencies]
 python = "^3.9"
 pandas = "^2.0.2"
-statsmodels = ">0.14.0"
 numpy = "^1.25.0"
 patsy = ">0.5.0"
 polars = ">0.18.3"
@@ -23,6 +22,7 @@ mkdocs = "^1.4.3"
 mkdocs-material = "^9.1.17"
 mkautodoc = ">=0.2.0"
 matplotlib = "^3.7.1"
+statsmodels = ">0.14.0"
 typing-extensions = "^4.7.0"
 pytest-xdist = "^3.3.1"
 bandit = "^1.7.5"