Return VarName type from get_var_name function

Author: michaelosthege

pymc/model/core.py

@@ -66,7 +66,7 @@
 from pymc.initial_point import make_initial_point_fn
 from pymc.logprob.basic import transformed_conditional_logp
 from pymc.logprob.utils import ParameterValueError
-from pymc.model_graph import VarName, model_to_graphviz
+from pymc.model_graph import model_to_graphviz
 from pymc.pytensorf import (
     PointFunc,
     SeedSequenceSeed,
@@ -80,6 +80,7 @@
 )
 from pymc.util import (
     UNSET,
+    VarName,
     WithMemoization,
     _add_future_warning_tag,
     get_transformed_name,
@@ -2061,7 +2062,7 @@ def compile_fn(
     )
 
 
-def Point(*args, filter_model_vars=False, **kwargs) -> Dict[str, np.ndarray]:
+def Point(*args, filter_model_vars=False, **kwargs) -> Dict[VarName, np.ndarray]:
     """Build a point. Uses same args as dict() does.
     Filters out variables not in the model. All keys are strings.
 

pymc/model_graph.py

@@ -14,7 +14,7 @@
 import warnings
 
 from collections import defaultdict
-from typing import Dict, Iterable, List, NewType, Optional, Sequence, Set
+from typing import Dict, Iterable, List, Optional, Sequence, Set
 
 from pytensor import function
 from pytensor.compile.sharedvalue import SharedVariable
@@ -28,10 +28,7 @@
 
 import pymc as pm
 
-from pymc.util import get_default_varnames, get_var_name
-
-VarName = NewType("VarName", str)
-
+from pymc.util import VarName, get_default_varnames, get_var_name
 
 __all__ = (
     "ModelGraph",
@@ -113,7 +110,7 @@ def vars_to_plot(self, var_names: Optional[Iterable[VarName]] = None) -> List[Va
                 selected_ancestors.add(self.model.rvs_to_values[var])
 
         # ordering of self._all_var_names is important
-        return [VarName(get_var_name(var)) for var in selected_ancestors]
+        return [get_var_name(var) for var in selected_ancestors]
 
     def make_compute_graph(
         self, var_names: Optional[Iterable[VarName]] = None

pymc/util.py

@@ -15,7 +15,7 @@
 import functools
 import warnings
 
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Dict, List, NewType, Optional, Sequence, Tuple, Union, cast
 
 import arviz
 import cloudpickle
@@ -29,6 +29,8 @@
 
 from pymc.exceptions import BlockModelAccessError
 
+VarName = NewType("VarName", str)
+
 
 class _UnsetType:
     """Type for the `UNSET` object to make it look nice in `help(...)` outputs."""
@@ -207,9 +209,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) -> str:
+def get_var_name(var) -> VarName:
     """Get an appropriate, plain variable name for a variable."""
-    return str(getattr(var, "name", var))
+    return VarName(str(getattr(var, "name", var)))
 
 
 def get_transformed(z):