Repository: matthewwardrop/formulaic

Issue:
Add required variables to the `Formula` class
I would like to be able to do something like the following. Appologies I am struggling to articulate what I want but effectively I want the following.

Say I have the following formula.
`apps ~ prior_apps + I(prior_apps^2) + factor + I(prior_apps:factor)`

I am wondering if it is possible to get extract out the rhs terms from the formula. By terms I mean `['prior_apps','factor']`

I have tried doing the following.

```
formula_parser = formulaic.parser.DefaultFormulaParser()
tokens = formula_parser.get_tokens(formula_str)
tokens = [t for t in tokens]
```
but that gets me the individual parts of the string and not the terms.

I feel like it should be possible?

Hints from issue comments:
@matthewwardrop: Hi @timpiperseek ,

Does something like the following work?

```
from formulaic import Formula
f = Formula("apps ~ prior_apps + I(prior_apps**2) + factor + prior_apps:factor")
set(
    factor
    for term in f.rhs
    for factor in term.factors
)
# This would output all the factors: {1, I(prior_apps ** 2), factor, prior_apps}
```
 (Note that interaction terms should not be enclosed in "I(...)", since that is a Python function call).

If you need to, you could parse the AST represented by the  non-lookup factors (e.g. `I(prior_apps ** 2)`) to extract the variables used; `prior_apps` here. 

If you are actually just looking for the terms, you can do: `list(f.rhs) == [1, prior_apps, I(prior_apps ** 2), factor, prior_apps:factor]`.

Does that help?@timpiperseek: yeah that is really close to what I am after.

what do you mean by 

> If you need to, you could parse the AST represented by the non-lookup factors (e.g. I(prior_apps ** 2)) to extract the variables used; prior_apps here.

because ideally it would also identify that prior_apps**2 is the same underlying metric as prior_apps.@matthewwardrop: Ah... Using some internal utility functions you can do:

```
from formulaic import Formula
from formulaic.utils.variables import get_expression_variables
f = Formula("apps ~ prior_apps + I(prior_apps**2) + factor + prior_apps:factor")
set(
    variable
    for term in f.rhs
    for factor in term.factors
    for variable in get_expression_variables(factor.expr, {})
    if "value" in variable.roles
)
# Outputs: {'factor', 'prior_apps'}
```

Note that `get_expression_variables` parses the AST associated with the python expression, which is used internally to keep track of which variables have been used when generating the model matrix.
@timpiperseek: Oh that is absolutely awesome, thank you. @matthewwardrop: I'll consider adding this directly to the formula class as something like `.required_variables`.@mayer79: This would indeed be very handy, thx.@matthewwardrop: Finally got around to this. This will be available in Formulaic 1.1, which will be released shortly@mayer79: Fantastic!

Solution patch:
diff --git a/formulaic/formula.py b/formulaic/formula.py
index 6dd450cc..422267fe 100644
--- a/formulaic/formula.py
+++ b/formulaic/formula.py
@@ -10,6 +10,7 @@
 from .parser import DefaultFormulaParser
 from .parser.types import FormulaParser, OrderedSet, Structured, Term
 from .utils.calculus import differentiate_term
+from .utils.variables import Variable, get_expression_variables
 
 FormulaSpec: TypeAlias = Union[
     str,
@@ -244,6 +245,44 @@ def get_model_matrix(
             data, context=context, drop_rows=drop_rows
         )
 
+    @property
+    def required_variables(self) -> Set[Variable]:
+        """
+        The set of variables required in the data order to materialize this
+        formula.
+
+        Attempts are made to restrict these variables only to those expected in
+        the data, and not, for example, those associated with transforms and/or
+        values present in the evaluation namespace by default (e.g. `y ~ C(x)`
+        would include only `y` and `x`). This may not always be possible for
+        more advanced formulae that insert constants into the formula via the
+        evaluation context rather than the data context.
+        """
+
+        variables: List[Variable] = []
+
+        # Recurse through formula to collect all variables
+        self._map(
+            lambda terms: variables.extend(
+                variable
+                for term in terms
+                for factor in term.factors
+                for variable in get_expression_variables(factor.expr, {})
+                if "value" in variable.roles
+            )
+        )
+
+        # Filter out constants like `contr` that are already present in the
+        # TRANSFORMS namespace.
+        from formulaic.transforms import TRANSFORMS
+
+        return set(
+            filter(
+                lambda variable: variable.split(".", 1)[0] not in TRANSFORMS,
+                Variable.union(variables),
+            )
+        )
+
     def differentiate(  # pylint: disable=redefined-builtin
         self,
         *wrt: str,


Test patch:
diff --git a/tests/test_formula.py b/tests/test_formula.py
index 607c6c42..2f5d10db 100644
--- a/tests/test_formula.py
+++ b/tests/test_formula.py
@@ -200,3 +200,10 @@ def test_pickleable(self, formula_exprs):
         formula = pickle.load(o)
         assert formula.lhs.root == ["a"]
         assert formula.rhs.root == ["1", "b"]
+
+    def test_required_variables(self):
+        assert Formula("a + b").required_variables == {"a", "b"}
+        assert Formula("a + b + a:b").required_variables == {"a", "b"}
+        assert Formula("a + C(b)").required_variables == {"a", "b"}
+        assert Formula("a + C(b, levels=['b1', 'b2'])").required_variables == {"a", "b"}
+        assert Formula("a + C(b, contr.Treatment)").required_variables == {"a", "b"}