y0-causal-inference · cthoyt · Jun 20, 2025 · Jun 20, 2025
diff --git a/pyproject.toml b/pyproject.toml
@@ -67,9 +67,9 @@ dependencies = [
     "scikit-learn",
     "statsmodels",
     "scipy",
+    "typing-extensions",
 ]
 
-
 # see https://peps.python.org/pep-0735/ and https://docs.astral.sh/uv/concepts/dependencies/#dependency-groups
 [dependency-groups]
 tests = [

diff --git a/src/y0/algorithm/estimation/estimators.py b/src/y0/algorithm/estimation/estimators.py
@@ -65,16 +65,20 @@ def get_primal_ipw_point_estimate(
     treatment: Variable,
     treatment_value: int | float,
     outcome: Variable,
-) -> float:
+) -> float | Any:
     """Estimate the counterfactual mean E[Y(t)] with the Primal IPW estimator on p-fixable graphs."""
+    # TODO: This function currently returns type Any to conform with mypy requirements.
+    #       That's not the best reason to make a return type less restrictive.
+    #       Consider warning the user if the type of the return value cannot be
+    #       coerced to a float. See Issue #294. -callahanr
     beta_primal = get_beta_primal(
         data=data,
         graph=graph,
         treatment=treatment,
         treatment_value=treatment_value,
         outcome=outcome,
     )
-    return cast(float, np.mean(beta_primal).item())
+    return np.mean(beta_primal).item()
 
 
 def get_beta_primal(

diff --git a/src/y0/dsl.py b/src/y0/dsl.py
@@ -1564,7 +1564,9 @@ def _iter_variables(self) -> Iterable[Variable]:
         yield from self.domain
 
 
-Q = QFactor
+# We need to declare a type for this alias to avoid false MyPy errors.
+# See https://github.com/python/mypy/issues/7568
+Q: type[QFactor] = QFactor
 
 AA = Variable("AA")
 A, B, C, D, E, F, G, M, R, S, T, U, W, X, Y, Z = map(Variable, "ABCDEFGMRSTUWXYZ")

diff --git a/tests/data/generate.py b/tests/data/generate.py
@@ -19,7 +19,7 @@
 from tqdm import tqdm, trange
 
 from y0.algorithm.estimation import estimate_ace
-from y0.examples import examples
+from y0.examples import examples  # type: ignore[attr-defined]
 
 warnings.simplefilter(action="ignore", category=FutureWarning)
 

diff --git a/tests/test_algorithm/test_cg.py b/tests/test_algorithm/test_cg.py
@@ -27,7 +27,7 @@
     value_of_self_intervention,
 )
 from y0.dsl import A, B, D, Event, W, X, Y, Z
-from y0.examples import (
+from y0.examples import (  # type: ignore[attr-defined]
     figure_9a,
     figure_9b,
     figure_9c,
@@ -47,11 +47,11 @@ class TestCounterfactualGraph(cases.GraphTestCase):
     def test_world(self):
         """Test that a world contains an intervention."""
         with self.assertRaises(TypeError):
-            input_world1: World = World([-x])
+            input_world1: World = World([-x])  # type: ignore[annotation-unchecked]
             3 in input_world1  # noqa
 
         with self.assertRaises(TypeError):
-            input_world1: World = World([3])
+            input_world1: World = World([3])  # type: ignore[annotation-unchecked]
             3 in input_world1  # noqa
 
         input_world2 = World([-x])
@@ -94,7 +94,7 @@ def test_has_same_function(self):
 
     def test_nodes_attain_same_value(self):
         """Test that two variables attain the same value."""
-        event: Event = {D: -d}
+        event: Event = {D: -d}  # type: ignore[annotation-unchecked]
         self.assertTrue(nodes_attain_same_value(figure_11a.graph, event, D, D @ -d))
         self.assertTrue(nodes_attain_same_value(figure_11a.graph, event, D @ -d, D))
         self.assertTrue(
@@ -143,7 +143,7 @@ def test_has_same_confounders(self):
     def test_parents_attain_same_values(self):
         """Test that the parents of two nodes attain the same value."""
         graph = figure_9b.graph
-        event: Event = {Y @ -x: -y, D: -d, Z @ -d: -z, X: +x}
+        event: Event = {Y @ -x: -y, D: -d, Z @ -d: -z, X: +x}  # type: ignore[annotation-unchecked]
         self.assertTrue(parents_attain_same_values(figure_11a.graph, event, Z, Z @ -d))
         self.assertTrue(parents_attain_same_values(figure_11a.graph, event, Z, Z @ -x))
         self.assertTrue(parents_attain_same_values(figure_11a.graph, event, Z @ -d, Z @ -x))
@@ -544,7 +544,7 @@ def test_get_directed_edges(self):
 
     def test_is_pw_equivalent(self):
         """Test that two nodes in a parallel world graph are the same (lemma 24)."""
-        event: Event = {Y @ -x: -y, D: -d, Z @ -d: -z, X: +x}
+        event: Event = {Y @ -x: -y, D: -d, Z @ -d: -z, X: +x}  # type: ignore[annotation-unchecked]
         self.assertTrue(is_pw_equivalent(figure_9b.graph, event, D @ -X, D))
         self.assertTrue(is_pw_equivalent(figure_9b.graph, event, X @ -D, X))
         self.assertTrue(is_pw_equivalent(figure_11a.graph, event, Z, Z @ -X))

diff --git a/tests/test_algorithm/test_conditional_independencies.py b/tests/test_algorithm/test_conditional_independencies.py
@@ -11,7 +11,7 @@
     get_conditional_independencies,
 )
 from y0.dsl import AA, B, C, D, E, F, G, Variable, X, Y
-from y0.examples import (
+from y0.examples import (  # type: ignore[attr-defined]
     Example,
     d_separation_example,
     examples,

diff --git a/tests/test_algorithm/test_estimation.py b/tests/test_algorithm/test_estimation.py
@@ -7,7 +7,7 @@
 from y0.algorithm.estimation import ananke_average_causal_effect, df_covers_graph
 from y0.algorithm.estimation.estimators import get_primal_ipw_ace, get_state_space_map
 from y0.dsl import Variable
-from y0.examples import examples, frontdoor, napkin, napkin_example
+from y0.examples import examples, frontdoor, napkin, napkin_example  # type: ignore[attr-defined]
 from y0.graph import ANANKE_REQUIRED, is_p_fixable
 
 TOLERANCE = 0.1

diff --git a/tests/test_algorithm/test_falsification.py b/tests/test_algorithm/test_falsification.py
@@ -8,7 +8,7 @@
 
 from y0.algorithm.conditional_independencies import get_conditional_independencies
 from y0.algorithm.falsification import get_falsifications, get_graph_falsifications
-from y0.examples import asia_example, frontdoor_example
+from y0.examples import asia_example, frontdoor_example  # type: ignore[attr-defined]
 from y0.struct import get_conditional_independence_tests
 
 

diff --git a/tests/test_algorithm/test_id_alg.py b/tests/test_algorithm/test_id_alg.py
@@ -37,7 +37,7 @@
     Z,
     get_outcomes_and_treatments,
 )
-from y0.examples import (
+from y0.examples import (  # type: ignore[attr-defined]
     figure_6a,
     line_1_example,
     line_2_example,

diff --git a/tests/test_algorithm/test_id_star.py b/tests/test_algorithm/test_id_star.py
@@ -28,7 +28,7 @@
     idc_star,
 )
 from y0.dsl import D, One, P, Sum, W, X, Y, Z, Zero
-from y0.examples import (
+from y0.examples import (  # type: ignore[attr-defined]
     figure_9a,
     figure_9c,
     figure_9d,

diff --git a/tests/test_algorithm/test_original_id_star.py b/tests/test_algorithm/test_original_id_star.py
@@ -27,7 +27,7 @@
     idc_star,
 )
 from y0.dsl import D, One, P, Sum, W, X, Y, Z, Zero
-from y0.examples import (
+from y0.examples import (  # type: ignore[attr-defined]
     figure_9a,
     figure_9c,
     figure_9d,

diff --git a/tests/test_algorithm/test_simplify_latent.py b/tests/test_algorithm/test_simplify_latent.py
@@ -15,7 +15,7 @@
 )
 from y0.algorithm.taheri_design import taheri_design_dag
 from y0.dsl import U1, U2, U3, Y1, Y2, Y3, U, Variable, W
-from y0.examples import igf_example
+from y0.examples import igf_example  # type: ignore[attr-defined]
 from y0.graph import set_latent
 
 X1, X2, X3 = map(Variable, ["X1", "X2", "X3"])

diff --git a/tests/test_algorithm/test_tian_pearl_identify.py b/tests/test_algorithm/test_tian_pearl_identify.py
@@ -675,13 +675,15 @@ class TestComputeCFactorMarginalizingOverTopologicalSuccessors(cases.GraphTestCa
     expected_result_2 = Fraction(expected_result_2_num, expected_result_2_den)
 
     # Same thing, but with population probabilities
+    # @cthoyt mypy throws an error for each appearance of PP[] below:
+    # error: Value of type "type[PopulationProbabilityBuilderType]" is not indexable  [index]
     result_piece_pp = Product.safe(
         [
-            PP[Pi1](W1),
-            PP[Pi1](W3 | W1),
-            PP[Pi1](W2 | (W3, W1)),
-            PP[Pi1](X | (W1, W3, W2, W4)),
-            PP[Pi1](Y | (W1, W3, W2, W4, X)),
+            PP[Pi1](W1),  # type: ignore[index]
+            PP[Pi1](W3 | W1),  # type: ignore[index]
+            PP[Pi1](W2 | (W3, W1)),  # type: ignore[index]
+            PP[Pi1](X | (W1, W3, W2, W4)),  # type: ignore[index]
+            PP[Pi1](Y | (W1, W3, W2, W4, X)),  # type: ignore[index]
         ]
     )
     expected_result_1_part_1_pp = Fraction(

diff --git a/tests/test_algorithm/test_transport.py b/tests/test_algorithm/test_transport.py
@@ -43,7 +43,9 @@
     Z,
     Zero,
 )
-from y0.examples import tikka_trso_figure_8_graph as tikka_trso_figure_8
+from y0.examples import (  # type: ignore[attr-defined]
+    tikka_trso_figure_8_graph as tikka_trso_figure_8,
+)
 from y0.graph import NxMixedGraph
 from y0.mutate import canonicalize, fraction_expand
 

diff --git a/tests/test_causaleffect.py b/tests/test_causaleffect.py
@@ -5,7 +5,7 @@
 import pyparsing
 
 from y0.dsl import P, Q, Sum, Variable
-from y0.examples import examples, verma_1
+from y0.examples import examples, verma_1  # type: ignore[attr-defined]
 
 try:
     from rpy2 import robjects
@@ -38,7 +38,11 @@ def setUpClass(cls) -> None:
         except Exception as e:
             raise unittest.SkipTest(f"R packages not properly installed.\n\n{e}") from None
 
-        cls.assertIsNotNone(robjects, msg="make sure this was imported correctly.")
+        cls.assertIsNotNone(
+            unittest.TestCase(), robjects, msg="make sure this was imported correctly."
+        )
+        # callahanr: @cthoyt Please take a look at this, adding the first parameter is a fix to correct a mypy error:
+        #            error: Missing positional argument "obj" in call to "assertIsNotNone" of "TestCase"  [call-arg]
 
     def test_verma_constraint(self):
         """Test getting the single Verma constraint from the Figure 1A graph."""

diff --git a/tests/test_graph.py b/tests/test_graph.py
@@ -7,7 +7,13 @@
 from pgmpy.models import DiscreteBayesianNetwork
 
 from y0.dsl import V1, V2, V3, V4, A, B, C, D, M, Variable, X, Y, Z
-from y0.examples import SARS_SMALL_GRAPH, Example, examples, napkin, verma_1
+from y0.examples import (  # type: ignore[attr-defined]
+    SARS_SMALL_GRAPH,
+    Example,
+    examples,
+    napkin,
+    verma_1,
+)
 from y0.graph import (
     ANANKE_AVAILABLE,
     ANANKE_REQUIRED,

diff --git a/tox.ini b/tox.ini
@@ -126,7 +126,7 @@ dependency_groups =
     typing
 commands =
     mypy --ignore-missing-imports --strict src/
-    mypy --ignore-missing-imports src/ tests/test_dsl.py tests/test_algorithm/test_counterfactual_transportability.py
+    mypy --ignore-missing-imports src/ tests/
 
 [testenv:docs-lint]
 skip_install = true