Replace VarName usage with str in type hints and conversion

Author: osyuksel

pymc/model/core.py

@@ -1939,7 +1939,7 @@ def debug_parameters(rv):
     def to_graphviz(
         self,
         *,
-        var_names: Iterable[VarName] | None = None,
+        var_names: Iterable[str] | None = None,
         formatting: str = "plain",
         save: str | None = None,
         figsize: tuple[int, int] | None = None,
@@ -2143,7 +2143,7 @@ def compile_fn(
     )
 
 
-def Point(*args, filter_model_vars=False, **kwargs) -> dict[VarName, np.ndarray]:
+def Point(*args, filter_model_vars=False, **kwargs) -> dict[str, np.ndarray]:
     """Build a point.
 
     Uses same args as dict() does.

pymc/model_graph.py

@@ -170,7 +170,7 @@ def default_data(var: TensorVariable) -> GraphvizNodeKwargs:
     }
 
 
-def get_node_type(var_name: VarName, model) -> NodeType:
+def get_node_type(var_name: str, model) -> NodeType:
     """Return the node type of the variable in the model."""
     v = model[var_name]
 
@@ -239,7 +239,7 @@ def __init__(self, model):
         self._all_vars = {model[var_name] for var_name in self._all_var_names}
         self.var_list = self.model.named_vars.values()
 
-    def get_parent_names(self, var: TensorVariable) -> set[VarName]:
+    def get_parent_names(self, var: TensorVariable) -> set[str]:
         if var.owner is None:
             return set()
 
@@ -258,12 +258,12 @@ def _expand(x):
             return x.owner.inputs
 
         return {
-            cast(VarName, ancestor.name)  # type: ignore[union-attr]
+            cast(str, ancestor.name)  # type: ignore[union-attr]
             for ancestor in walk(nodes=var.owner.inputs, expand=_expand)
             if ancestor in named_vars
         }
 
-    def vars_to_plot(self, var_names: Iterable[VarName] | None = None) -> list[VarName]:
+    def vars_to_plot(self, var_names: Iterable[str] | None = None) -> list[str]:
         if var_names is None:
             return self._all_var_names
 
@@ -294,12 +294,12 @@ def vars_to_plot(self, var_names: Iterable[VarName] | None = None) -> list[VarNa
         return [get_var_name(var) for var in selected_ancestors]
 
     def make_compute_graph(
-        self, var_names: Iterable[VarName] | None = None
-    ) -> dict[VarName, set[VarName]]:
+        self, var_names: Iterable[str] | None = None
+    ) -> dict[str, set[str]]:
         """Get map of var_name -> set(input var names) for the model."""
         model = self.model
         named_vars = self._all_vars
-        input_map: dict[VarName, set[VarName]] = defaultdict(set)
+        input_map: dict[str, set[str]] = defaultdict(set)
 
         var_names_to_plot = self.vars_to_plot(var_names)
         for var_name in var_names_to_plot:
@@ -316,15 +316,15 @@ def make_compute_graph(
                 for ancestor in ancestors([obs_var]):
                     if ancestor not in named_vars:
                         continue
-                    obs_name = cast(VarName, ancestor.name)
+                    obs_name = cast(str, ancestor.name)
                     input_map[var_name].discard(obs_name)
                     input_map[obs_name].add(var_name)
 
         return input_map
 
     def get_plates(
         self,
-        var_names: Iterable[VarName] | None = None,
+        var_names: Iterable[str] | None = None,
     ) -> list[Plate]:
         """Rough but surprisingly accurate plate detection.
 
@@ -386,8 +386,8 @@ def get_plates(
 
     def edges(
         self,
-        var_names: Iterable[VarName] | None = None,
-    ) -> list[tuple[VarName, VarName]]:
+        var_names: Iterable[str] | None = None,
+    ) -> list[tuple[str, str]]:
         """Get edges between the variables in the model.
 
         Parameters
@@ -402,7 +402,7 @@ def edges(
 
         """
         return [
-            (VarName(child.replace(":", "&")), VarName(parent.replace(":", "&")))
+            (str(child.replace(":", "&")), str(parent.replace(":", "&")))
             for child, parents in self.make_compute_graph(var_names=var_names).items()
             for parent in parents
         ]
@@ -419,7 +419,7 @@ def nodes(self, plates: list[Plate] | None = None) -> list[NodeInfo]:
 def make_graph(
     name: str,
     plates: list[Plate],
-    edges: list[tuple[VarName, VarName]],
+    edges: list[tuple[str, str]],
     formatting: str = "plain",
     save=None,
     figsize=None,
@@ -493,7 +493,7 @@ def make_graph(
 def make_networkx(
     name: str,
     plates: list[Plate],
-    edges: list[tuple[VarName, VarName]],
+    edges: list[tuple[str, str]],
     formatting: str = "plain",
     node_formatters: NodeTypeFormatterMapping | None = None,
     create_plate_label: PlateLabelFunc = create_plate_label_with_dim_length,
@@ -563,7 +563,7 @@ def make_networkx(
 def model_to_networkx(
     model=None,
     *,
-    var_names: Iterable[VarName] | None = None,
+    var_names: Iterable[str] | None = None,
     formatting: str = "plain",
     node_formatters: NodeTypeFormatterMapping | None = None,
     include_dim_lengths: bool = True,
@@ -657,7 +657,7 @@ def model_to_networkx(
 def model_to_graphviz(
     model=None,
     *,
-    var_names: Iterable[VarName] | None = None,
+    var_names: Iterable[str] | None = None,
     formatting: str = "plain",
     save: str | None = None,
     figsize: tuple[int, int] | None = None,

pymc/util.py

@@ -214,9 +214,9 @@ def get_default_varnames(var_iterator, include_transformed):
         return [var for var in var_iterator if not is_transformed_name(get_var_name(var))]
 
 
-def get_var_name(var) -> VarName:
+def get_var_name(var) -> str:
     """Get an appropriate, plain variable name for a variable."""
-    return VarName(str(getattr(var, "name", var)))
+    return str(getattr(var, "name", var))
 
 
 def get_transformed(z):