Make paths() function match docs

Author: mdko

pyrtl/analysis.py

@@ -410,10 +410,10 @@ def extract_area_delay_from_yosys_output(yosys_output):
 def paths(src=None, dst=None, dst_nets=None, block=None):
     """ Get the list of paths from src to dst.
 
-    :param WireVector src: source wire(s) from which to trace your paths;
-        if None, will get paths from all Inputs
-    :param WireVector dst: destination wire(s) to which to trace your paths
-        if None, will get paths to all Outputs
+    :param Union[WireVector, Iterable[WireVector]] src: source wire(s) from which to
+        trace your paths; if None, will get paths from all Inputs
+    :param Union[WireVector, Iterable[WireVector]] dst: destination wire(s) to which to
+        trace your paths; if None, will get paths to all Outputs
     :param {WireVector: {LogicNet}} dst_nets: map from wire to set of nets where the
         wire is an argument; will compute it internally if not given via a
         call to pyrtl.net_connections()
@@ -448,10 +448,21 @@ def paths(src=None, dst=None, dst_nets=None, block=None):
         for output in block.wirevector_subset(cls=Output):
             dst_nets.pop(output, None)
 
-    src = block.wirevector_subset(cls=Input) if src is None else {src}
-    dst = block.wirevector_subset(cls=Output) if dst is None else {dst}
+    if src is None:
+        src = block.wirevector_subset(cls=Input)
+    elif isinstance(src, WireVector):
+        src = {src}
+    else:
+        src = set(src)
+
+    if dst is None:
+        dst = block.wirevector_subset(cls=Output)
+    elif isinstance(dst, WireVector):
+        dst = {dst}
+    else:
+        dst = set(dst)
 
-    def paths_src_dst(src, dst, block=None):
+    def paths_src_dst(src, dst):
         paths = []
 
         # Use DFS to get the paths [each a list of nets] from src wire to dst wire

tests/test_analysis.py

@@ -162,6 +162,25 @@ def test_two_paths_to_two_outputs(self):
         path_to_o2 = paths_from_a[o2]
         self.assertEqual(len(path_to_o2), 1)
 
+    def test_subset_of_all_paths(self):
+        i, j, k = pyrtl.input_list('i/2 j/3 k/4')
+        o, p = pyrtl.Output(), pyrtl.Output()
+        o <<= i & j
+        p <<= k - i
+
+        # Make sure passing in both set and list works
+        paths = pyrtl.paths([i, k], {o})
+        paths_from_i = paths[i]
+        self.assertNotIn(p, paths_from_i)  # Because p was not provided as target output
+        self.assertEqual(len(paths_from_i[o]), 1)  # One path from i to o
+        self.assertEqual(paths_from_i[o][0][0].op, 'c')
+        self.assertEqual(paths_from_i[o][0][1].op, '&')
+        self.assertEqual(paths_from_i[o][0][2].op, 'w')
+
+        paths_from_k = paths[k]
+        self.assertNotIn(p, paths_from_k)  # Because p was not provided as target output
+        self.assertEqual(len(paths_from_k[o]), 0)  # 0 paths from k to o
+
     def test_all_paths(self):
         a, b, c = pyrtl.input_list('a/2 b/4 c/1')
         o, p = pyrtl.output_list('o/4 p/2')