Lazy pydot import

Author: ricardoV94

pytensor/printing.py

@@ -26,39 +26,6 @@
 
 IDTypesType = Literal["id", "int", "CHAR", "auto", ""]
 
-pydot_imported = False
-pydot_imported_msg = ""
-try:
-    # pydot-ng is a fork of pydot that is better maintained
-    import pydot_ng as pd
-
-    if pd.find_graphviz():
-        pydot_imported = True
-    else:
-        pydot_imported_msg = "pydot-ng can't find graphviz. Install graphviz."
-except ImportError:
-    try:
-        # fall back on pydot if necessary
-        import pydot as pd
-
-        if hasattr(pd, "find_graphviz"):
-            if pd.find_graphviz():
-                pydot_imported = True
-            else:
-                pydot_imported_msg = "pydot can't find graphviz"
-        else:
-            pd.Dot.create(pd.Dot())
-            pydot_imported = True
-    except ImportError:
-        # tests should not fail on optional dependency
-        pydot_imported_msg = (
-            "Install the python package pydot or pydot-ng. Install graphviz."
-        )
-    except Exception as e:
-        pydot_imported_msg = "An error happened while importing/trying pydot: "
-        pydot_imported_msg += str(e.args)
-
-
 _logger = logging.getLogger("pytensor.printing")
 VALID_ASSOC = {"left", "right", "either"}
 
@@ -1288,6 +1255,45 @@ def pydotprint(
         scan separately after the top level debugprint output.
 
     """
+    pydot_imported = False
+    pydot_imported_msg = ""
+    try:
+        # pydot-ng is a fork of pydot that is better maintained
+        import pydot_ng as pd
+
+        if pd.find_graphviz():
+            pydot_imported = True
+        else:
+            pydot_imported_msg = "pydot-ng can't find graphviz. Install graphviz."
+    except ImportError:
+        try:
+            # fall back on pydot if necessary
+            import pydot as pd
+
+            if hasattr(pd, "find_graphviz"):
+                if pd.find_graphviz():
+                    pydot_imported = True
+                else:
+                    pydot_imported_msg = "pydot can't find graphviz"
+            else:
+                pd.Dot.create(pd.Dot())
+                pydot_imported = True
+        except ImportError:
+            # tests should not fail on optional dependency
+            pydot_imported_msg = (
+                "Install the python package pydot or pydot-ng. Install graphviz."
+            )
+        except Exception as e:
+            pydot_imported_msg = "An error happened while importing/trying pydot: "
+            pydot_imported_msg += str(e.args)
+
+    if not pydot_imported:
+        raise RuntimeError(
+            "Failed to import pydot. You must install graphviz "
+            "and either pydot or pydot-ng for "
+            f"`pydotprint` to work:\n {pydot_imported_msg}",
+        )
+
     from pytensor.scan.op import Scan
 
     if colorCodes is None:
@@ -1320,12 +1326,6 @@ def pydotprint(
         outputs = fct.outputs
         topo = fct.toposort()
         fgraph = fct
-    if not pydot_imported:
-        raise RuntimeError(
-            "Failed to import pydot. You must install graphviz "
-            "and either pydot or pydot-ng for "
-            f"`pydotprint` to work:\n {pydot_imported_msg}",
-        )
 
     g = pd.Dot()