misc: Cleanup + flake8 fixes

Author: enwask

devito/ir/iet/visitors.py

@@ -68,27 +68,21 @@ class LazyVisitor(GenericVisitor, Generic[ResultType]):
     A generic visitor that lazily yields results instead of flattening results
     from children at every step.
 
-    Subclass-defined visit methods (and default_retval) should be generators.
+    Subclass-defined visit methods should be generators.
     """
 
-    @classmethod
-    def default_retval(cls) -> Iterator[Any]:
-        yield from ()
-
     def lookup_method(self, instance) -> Callable[..., Iterator[Any]]:
         return super().lookup_method(instance)
 
     def _visit(self, o, *args, **kwargs) -> Iterator[Any]:
-        """Visit `o`."""
         meth = self.lookup_method(o)
         yield from meth(o, *args, **kwargs)
 
     def _post_visit(self, ret: Iterator[Any]) -> ResultType:
-        """Postprocess the visitor output before returning it to the caller."""
         return list(ret)
 
     def visit_object(self, o: object, **kwargs) -> Iterator[Any]:
-        yield from self.default_retval()
+        yield from ()
 
     def visit_Node(self, o: Node, **kwargs) -> Iterator[Any]:
         yield from self._visit(o.children, **kwargs)
@@ -1053,6 +1047,7 @@ class FindSymbols(LazyVisitor[list[Any]]):
         - `defines-aliases`: Collect all defined objects and their aliases
     """
 
+    @staticmethod
     def _defines_aliases(n):
         for i in n.defines:
             f = i.function

devito/symbolics/search.py

@@ -17,25 +17,19 @@
 Expression = sympy.Basic | np.number | int | float
 
 
-class Set(set[Expression]):
-    @staticmethod
-    def wrap(obj: Expression) -> set[Expression]:
-        return {obj}
-
-
 class List(list[Expression]):
-    @staticmethod
-    def wrap(obj: Expression) -> list[Expression]:
-        return [obj]
+    """
+    A list that aliases `extend` to `update` to mirror the `set` interface.
+    """
 
     def update(self, obj: Iterable[Expression]) -> None:
         self.extend(obj)
 
 
 Mode = Literal['all', 'unique']
-modes: dict[Mode, type[List] | type[Set]] = {
+modes: dict[Mode, type[List] | type[set[Expression]]] = {
     'all': List,
-    'unique': Set
+    'unique': set
 }
 
 
@@ -97,7 +91,7 @@ def search(exprs: Expression | Iterable[Expression],
            query: type | Callable[[Any], bool],
            mode: Mode = 'unique',
            visit: Literal['dfs', 'bfs', 'bfs_first_hit'] = 'dfs',
-           deep: bool = False) -> List | Set:
+           deep: bool = False) -> List | set[Expression]:
     """Interface to Search."""
 
     assert mode in ('all', 'unique'), "Unknown mode"
@@ -118,7 +112,7 @@ def search(exprs: Expression | Iterable[Expression],
         _search = searcher.visit_preorder_first_hit
     else:
         raise ValueError(f"Unknown visit mode '{visit}'")
-    
+
     exprs = filter(lambda e: isinstance(e, sympy.Basic), as_tuple(exprs))
     found = modes[mode](chain(*map(_search, exprs)))