Fix stat_quantile to handle formula that use environment variables

Author: has2k1

plotnine/stats/stat.py

@@ -64,7 +64,7 @@ class stat(ABC, metaclass=Register):
 
     # Plot namespace, it gets its value when the plot is being
     # built.
-    environment: Environment | None = None
+    environment: Environment
 
     def __init__(
         self,

plotnine/stats/stat_quantile.py

@@ -63,6 +63,13 @@ def setup_params(self, data):
         if params["formula"] is None:
             params["formula"] = "y ~ x"
             warn("Formula not specified, using '{}'", PlotnineWarning)
+        else:
+            from patsy.eval import EvalEnvironment
+
+            params["eval_env"] = EvalEnvironment(
+                namespaces=self.environment.namespaces
+            )
+
         try:
             iter(params["quantiles"])
         except TypeError:
@@ -81,7 +88,11 @@ def quant_pred(q, data, params):
     """
     import statsmodels.formula.api as smf
 
-    mod = smf.quantreg(params["formula"], data)
+    mod = smf.quantreg(
+        params["formula"],
+        data,
+        eval_env=params.get("eval_env"),
+    )
     reg_res = mod.fit(q=q, **params["method_args"])
     out = pd.DataFrame(
         {

tests/test_geom_quantile.py

@@ -18,7 +18,18 @@ def test_lines():
         + geom_quantile(quantiles=[0.001, 0.5, 0.999], formula="y~x", size=2)
     )
 
+    # np.absolute tests the ability to pickup variables in the
+    # caller environment
+    p2 = (
+        ggplot(data, aes(x="x", y="y"))
+        + geom_point(alpha=0.5)
+        + geom_quantile(
+            quantiles=[0.001, 0.5, 0.999], formula="y~np.absolute(x)", size=2
+        )
+    )
+
     # Two (.001, .999) quantile lines should bound the points
     # from below and from above, and the .5 line should go
     # through middle (approximately).
     assert p == "lines"
+    assert p2 == "lines"