Force Graph to actually use frozensets

Because type annotations in python are a "I wish people would respect my
authority" kind of feature, we cannot guarantee that a `Graph` is
actually created with frozensets as attributes. That is the case all
throughout taskgraph where we're passing sets. By forcing those
attributes to be frozensets we can ensure that `Graph` is hashable which
will allow us to use `functools.lru_cache` to cache some expensive
functions in here.

Unfortunately, there's no easy way of transforming the inputs passed
into a dataclass, since in both `__init__` and `__post_init__`, the
structure is already frozen. By using a dataclass as the parent of
`Graph` instead of `Graph` itself, we can work around the issue while
keeping the frozenness of the dataclass post init without having to
resort to working around dataclass when transforming the attributes

Author: Eijebong

src/taskgraph/graph.py

@@ -8,7 +8,12 @@
 
 
 @dataclass(frozen=True)
-class Graph:
+class _Graph:
+    nodes: frozenset
+    edges: frozenset
+
+
+class Graph(_Graph):
     """Generic representation of a directed acyclic graph with labeled edges
     connecting the nodes. Graph operations are implemented in a functional
     manner, so the data structure is immutable.
@@ -23,8 +28,8 @@ class Graph:
     node `left` to node `right`..
     """
 
-    nodes: frozenset
-    edges: frozenset
+    def __init__(self, nodes, edges):
+        super().__init__(frozenset(nodes), frozenset(edges))
 
     def transitive_closure(self, nodes, reverse=False):
         """Return the transitive closure of <nodes>: the graph containing all
@@ -67,7 +72,7 @@ def transitive_closure(self, nodes, reverse=False):
             add_nodes = {(left if reverse else right) for (left, right, _) in add_edges}
             new_nodes = nodes | add_nodes
             new_edges = edges | add_edges
-        return Graph(new_nodes, new_edges)  # type: ignore
+        return Graph(new_nodes, new_edges)
 
     def _visit(self, reverse):
         queue = collections.deque(sorted(self.nodes))

src/taskgraph/morph.py

@@ -51,7 +51,7 @@ def amend_taskgraph(taskgraph, label_to_taskid, to_add):
         for depname, dep in task.dependencies.items():
             new_edges.add((task.task_id, dep, depname))
 
-    taskgraph = TaskGraph(new_tasks, Graph(set(new_tasks), new_edges))  # type: ignore
+    taskgraph = TaskGraph(new_tasks, Graph(frozenset(new_tasks), new_edges))
     return taskgraph, label_to_taskid
 
 

src/taskgraph/optimize/base.py

@@ -451,7 +451,9 @@ def get_subgraph(
         if left in tasks_by_taskid and right in tasks_by_taskid
     }
 
-    return TaskGraph(tasks_by_taskid, Graph(set(tasks_by_taskid), edges_by_taskid))  # type: ignore
+    return TaskGraph(
+        tasks_by_taskid, Graph(frozenset(tasks_by_taskid), edges_by_taskid)
+    )
 
 
 @register_strategy("never")

src/taskgraph/taskgraph.py

@@ -67,5 +67,5 @@ def from_json(cls, tasks_dict):
                 tasks[key].task_id = value["task_id"]
             for depname, dep in value["dependencies"].items():
                 edges.add((key, dep, depname))
-        task_graph = cls(tasks, Graph(set(tasks), edges))  # type: ignore
+        task_graph = cls(tasks, Graph(frozenset(tasks), edges))
         return tasks, task_graph