Split off tree into its own thing.

Author: rocky

pymathics/graph/__init__.py

@@ -6,6 +6,7 @@
 
 
 from pymathics.graph.__main__ import *
+from pymathics.graph.tree import *
 from pymathics.graph.graph_generators import *
 from pymathics.graph.version import __version__
 

pymathics/graph/__main__.py

@@ -32,15 +32,7 @@
     "PlotLabel": "Null",
 }
 
-DEFAULT_TREE_OPTIONS = {**DEFAULT_GRAPH_OPTIONS,
-                        **{"Directed" : "False",
-                           "GraphLayout": '"tree"'}}
-
-try:
-    import networkx as nx
-except ImportError:
-    nx = {}
-
+import networkx as nx
 
 def _circular_layout(G):
     return nx.drawing.circular_layout(G, scale=1.5)
@@ -453,12 +445,6 @@ def update_weights(self, evaluation):
         return weights
 
 
-class TreeGraph(Graph):
-    def __init__(self, G, **kwargs):
-        super(Graph, self).__init__()
-        self.G = G
-
-
 def _is_connected(G):
     if len(G) == 0:  # empty graph?
         return True
@@ -776,44 +762,6 @@ def apply_1(self, vertices, edges, evaluation, options):
         return _graph_from_list(edges.leaves, options=options, new_vertices=vertices.leaves)
 
 
-class TreeGraphAtom(AtomBuiltin):
-    """
-    >> TreeGraph[{1->2, 2->3, 3->1}]
-     = -Graph-
-
-    """
-
-    options = DEFAULT_TREE_OPTIONS
-
-    messages = {
-        "v": "Expected first parameter vertices to be a list of vertices",
-        "notree": "Graph is not a tree.",
-    }
-
-    def apply(self, rules, evaluation, options):
-        "TreeGraph[rules_List, OptionsPattern[%(name)s]]"
-        g = _graph_from_list(rules.leaves, options)
-        if not nx.is_tree(g.G):
-            evaluation.message(self.get_name(), "notree")
-
-        g.G.graph_layout = String("tree")
-        # Compute/check/set for root?
-        return g
-
-    def apply_1(self, vertices, edges, evaluation, options):
-        "TreeGraph[vertices_List, edges_List, OptionsPattern[%(name)s]]"
-        if not all(isinstance(v, Atom) for v in vertices.leaves):
-            evaluation.message(self.get_name(), "v")
-
-        g = _graph_from_list(edges.leaves, options=options, new_vertices=vertices.leaves)
-        if not nx.is_tree(g.G):
-            evaluation.message(self.get_name(), "notree")
-
-        g.G.graph_layout = String("tree")
-        # Compute/check/set for root?
-        return g
-
-
 class PathGraph(_NetworkXBuiltin):
     """
     >> PathGraph[{1, 2, 3}]

pymathics/graph/graph_generators.py

@@ -1,9 +1,11 @@
 from pymathics.graph.__main__ import (
-    DEFAULT_TREE_OPTIONS,
     Graph,
     _NetworkXBuiltin,
     nx,
 )
+
+from pymathics.graph.tree import DEFAULT_TREE_OPTIONS
+
 from mathics.core.expression import Integer, String
 from typing import Callable, Optional
 

pymathics/graph/tree.py

@@ -0,0 +1,56 @@
+import networkx as nx
+from pymathics.graph.__main__ import Graph, _graph_from_list, DEFAULT_GRAPH_OPTIONS
+from mathics.core.expression import String
+
+DEFAULT_TREE_OPTIONS = {
+    **DEFAULT_GRAPH_OPTIONS,
+    **{"Directed": "False", "GraphLayout": '"tree"'},
+}
+
+from mathics.builtin.base import Builtin, AtomBuiltin
+
+
+class TreeGraphAtom(AtomBuiltin):
+    """
+    >> TreeGraph[{1->2, 2->3, 3->1}]
+     = -Graph-
+
+    """
+
+    options = DEFAULT_TREE_OPTIONS
+
+    messages = {
+        "v": "Expected first parameter vertices to be a list of vertices",
+        "notree": "Graph is not a tree.",
+    }
+
+    def apply(self, rules, evaluation, options):
+        "TreeGraph[rules_List, OptionsPattern[%(name)s]]"
+        g = _graph_from_list(rules.leaves, options)
+        if not nx.is_tree(g.G):
+            evaluation.message(self.get_name(), "notree")
+
+        g.G.graph_layout = String("tree")
+        # Compute/check/set for root?
+        return g
+
+    def apply_1(self, vertices, edges, evaluation, options):
+        "TreeGraph[vertices_List, edges_List, OptionsPattern[%(name)s]]"
+        if not all(isinstance(v, Atom) for v in vertices.leaves):
+            evaluation.message(self.get_name(), "v")
+
+        g = _graph_from_list(
+            edges.leaves, options=options, new_vertices=vertices.leaves
+        )
+        if not nx.is_tree(g.G):
+            evaluation.message(self.get_name(), "notree")
+
+        g.G.graph_layout = String("tree")
+        # Compute/check/set for root?
+        return g
+
+
+class TreeGraph(Graph):
+    def __init__(self, G, **kwargs):
+        super(Graph, self).__init__()
+        self.G = G