Fix types

Author: charmoniumQ

probe_py/probe_py/analysis.py

@@ -1,7 +1,6 @@
 import collections
 from .ptypes import ProbeLog, HbGraph, OpQuad
 from .ops import CloneOp, WaitOp
-from .hb_graph import HbGraph
 
 
 def get_max_parallelism_latest(hb_graph: HbGraph, probe_log: ProbeLog) -> int:

probe_py/probe_py/dataflow_graph.py

@@ -8,6 +8,7 @@
 import typing
 import warnings
 import networkx
+import tqdm
 from . import graph_utils
 from .hb_graph_accesses import hb_graph_to_accesses
 from . import ops
@@ -84,26 +85,26 @@ def ensure_state(quad: ptypes.OpQuad, desired_state: PidState) -> ptypes.OpQuint
                 inode_to_paths[access.inode].add(access.path)
                 version = InodeVersionNode(access.inode, version_num)
                 next_version = InodeVersionNode(access.inode, version_num + 1)
-                ensure_state(access.op_node, PidState.READING if access.mode.is_side_effect_free() else PidState.WRITING)
+                ensure_state(access.op_node, PidState.READING if access.mode.is_side_effect_free else PidState.WRITING)
                 if (op_node := last_op_in_process.get(access.op_node.pid)) is None:
                     warnings.warn(ptypes.UnusualProbeLog(f"Can't find last node from process {access.op_node.pid}"))
                     continue
                 match access.mode:
-                    case AccessMode.WRITE:
-                        if access.phase == Phase.BEGIN:
+                    case ptypes.AccessMode.WRITE:
+                        if access.phase == ptypes.Phase.BEGIN:
                             dataflow_graph.add_edge(op_node, next_version)
                             dataflow_graph.add_edge(version, next_version)
-                    case AccessMode.TRUNCATE_WRITE:
-                        if access.phase == Phase.END:
+                    case ptypes.AccessMode.TRUNCATE_WRITE:
+                        if access.phase == ptypes.Phase.END:
                             dataflow_graph.add_edge(op_node, next_version)
-                    case AccessMode.READ_WRITE:
-                        if access.phase == Phase.BEGIN:
+                    case ptypes.AccessMode.READ_WRITE:
+                        if access.phase == ptypes.Phase.BEGIN:
                             dataflow_graph.add_edge(version, op_node)
-                        if access.phase == Phase.END:
+                        if access.phase == ptypes.Phase.END:
                             dataflow_graph.add_edge(op_node, next_version)
                             dataflow_graph.add_edge(version, next_version)
-                    case AccessMode.READ | AccessMode.EXEC | AccessMode.DLOPEN:
-                        if access.phase == Phase.BEGIN:
+                    case ptypes.AccessMode.READ | ptypes.AccessMode.EXEC | ptypes.AccessMode.DLOPEN:
+                        if access.phase == ptypes.Phase.BEGIN:
                             dataflow_graph.add_edge(version, op_node)
                     case _:
                         raise TypeError()
@@ -233,9 +234,9 @@ def label_nodes(
     ):
         data = dataflow_graph.nodes(data=True)[node]
         match node:
-            case hb_graph.OpNode():
+            case ptypes.OpQuad():
                 data["shape"] = "oval"
-                op = probe_log.get_op(node.pid, node.exec_no, node.pid.main_thread(), 0)
+                op = probe_log.get_op(ptypes.OpQuad(node.pid, node.exec_no, node.pid.main_thread(), 0))
                 if node.op_no == 0:
                     count[(node.pid, node.exec_no)] = 1
                     if node.exec_no != 0:

probe_py/probe_py/graph_utils.py

@@ -86,8 +86,7 @@ def map_nodes(
 ) -> networkx.DiGraph[_Node2]:
     dct = {node: function(node) for node in graph.nodes()}
     assert util.all_unique(dct.values()), util.duplicates(dct.values())
-    ret = networkx.relabel_nodes(graph, dct)
-    return typing.cast(networkx.DiGraph[_Node2], ret)
+    return networkx.relabel_nodes(graph, dct)
 
 
 def serialize_graph(
@@ -99,7 +98,7 @@ def serialize_graph(
     if name_mapper is None:
         name_mapper = typing.cast(
             typing.Callable[[_Node], str],
-            lambda node: graph.nodes[node].get("id", str(node)),
+            lambda node: graph.nodes(data=True)[node].get("id", str(node)),
         )
     graph2 = map_nodes(name_mapper, graph)
     pydot_graph = networkx.drawing.nx_pydot.to_pydot(graph2)
@@ -376,7 +375,9 @@ def add_edge(self, source: _Node, target: _Node) -> None:
                     self.dag_tc.add_edge(descendant_of_source, descendant_of_target)
 
 
-def get_faces(planar_graph: networkx.PlanarEmbedding[_Node]) -> frozenset[tuple[_Node, ...]]:
+def get_faces(
+        planar_graph: networkx.PlanarEmbedding[_Node], # type: ignore
+) -> frozenset[tuple[_Node, ...]]:
     faces = set()
     covered_half_edges = set()
     for half_edge in planar_graph.edges():
@@ -475,7 +476,7 @@ def topological_sort_depth_first(
 ) -> typing.Iterable[_Node]:
     """Topological sort that breaks ties by depth first, and then by lowest child score."""
     queue = util.PriorityQueue[_Node, tuple[int, int]](
-        (node, (typing.cast(int, dag.in_degree(node)), 0))
+        (node, (dag.in_degree(node), 0))
         for node in dag.nodes()
     )
     counter = 0

probe_py/probe_py/hb_graph_accesses.py

@@ -59,11 +59,11 @@ def openfd(
         match op_data:
             case ops.InitExecEpochOp():
                 if 0 not in proc_fd_to_fd[node.pid]:
-                    yield from openfd(0, ptypes.AccessMode.READ, False, node, op_data.stdin)
+                    yield from openfd(0, ptypes.AccessMode.READ, False, node, op_data.std_in)
                 if 1 not in proc_fd_to_fd[node.pid]:
-                    yield from openfd(1, ptypes.AccessMode.TRUNCATE_WRITE, False, node, op_data.stdout)
+                    yield from openfd(1, ptypes.AccessMode.TRUNCATE_WRITE, False, node, op_data.std_out)
                 if 2 not in proc_fd_to_fd[node.pid]:
-                    yield from openfd(2, ptypes.AccessMode.TRUNCATE_WRITE, False, node, op_data.stderr)
+                    yield from openfd(2, ptypes.AccessMode.TRUNCATE_WRITE, False, node, op_data.std_err)
             case ops.OpenOp():
                 if op_data.ferrno == 0:
                     mode = ptypes.AccessMode.from_open_flags(op_data.flags)

probe_py/probe_py/ptypes.py

@@ -9,6 +9,7 @@
 import socket
 import stat
 import typing
+import networkx
 import numpy
 from . import ops
 from . import consts