fix: ruff and type hints

Author: giospada

src/substrait/builders/plan.py

@@ -5,8 +5,8 @@
 See `examples/builder_example.py` for usage.
 """
 
-from typing import Iterable, Optional, Union, Callable
 
+from typing import Callable, Iterable, Optional, TypedDict, Union
 import substrait.gen.proto.algebra_pb2 as stalg
 from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
 import substrait.gen.proto.plan_pb2 as stp
@@ -20,16 +20,23 @@
 from substrait.type_inference import infer_plan_schema
 from substrait.utils import (
     merge_extension_declarations,
-    merge_extension_urns,
     merge_extension_uris,
+    merge_extension_urns,
 )
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 
 PlanOrUnbound = Union[stp.Plan, UnboundPlan]
 
+_ExtensionDict = TypedDict(
+    "_ExtensionDict",
+    {"extension_uris": list, "extension_urns": list, "extensions": list},
+)
+
 
-def _merge_extensions(*objs):
+def _merge_extensions(
+    *objs,
+) -> _ExtensionDict:
     """Merge extension URIs, URNs, and declarations from multiple plan/expression objects.
 
     During the URI -> URN migration period, we maintain both URI and URN references

src/substrait/derivation_expression.py

@@ -1,5 +1,7 @@
 from typing import Optional
-from antlr4 import InputStream, CommonTokenStream
+
+from antlr4 import CommonTokenStream, InputStream
+
 from substrait.gen.antlr.SubstraitTypeLexer import SubstraitTypeLexer
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
 from substrait.gen.proto.type_pb2 import Type
@@ -9,8 +11,9 @@ def _evaluate(x, values: dict):
     if isinstance(x, SubstraitTypeParser.BinaryExprContext):
         left = _evaluate(x.left, values)
         right = _evaluate(x.right, values)
-
-        if x.op.text == "+":
+        if x.op is None:
+            raise Exception("Undefined operator op")
+        elif x.op.text == "+":
             return left + right
         elif x.op.text == "-":
             return left - right
@@ -121,15 +124,19 @@ def _evaluate(x, values: dict):
                         nullability=nullability,
                     )
                 )
-            elif isinstance(parametrized_type, SubstraitTypeParser.PrecisionTimestampContext):
+            elif isinstance(
+                parametrized_type, SubstraitTypeParser.PrecisionTimestampContext
+            ):
                 precision = _evaluate(parametrized_type.precision, values)
                 return Type(
                     precision_timestamp=Type.PrecisionTimestamp(
                         precision=precision,
                         nullability=nullability,
                     )
                 )
-            elif isinstance(parametrized_type, SubstraitTypeParser.PrecisionTimestampTZContext):
+            elif isinstance(
+                parametrized_type, SubstraitTypeParser.PrecisionTimestampTZContext
+            ):
                 precision = _evaluate(parametrized_type.precision, values)
                 return Type(
                     precision_timestamp_tz=Type.PrecisionTimestampTZ(
@@ -144,27 +151,29 @@ def _evaluate(x, values: dict):
                     )
                 )
             elif isinstance(parametrized_type, SubstraitTypeParser.StructContext):
-                types = list(map(lambda x: _evaluate(x,values),parametrized_type.expr()))
+                types = list(
+                    map(lambda x: _evaluate(x, values), parametrized_type.expr())
+                )
                 return Type(
                     struct=Type.Struct(
                         types=types,
                         nullability=nullability,
                     )
                 )
             elif isinstance(parametrized_type, SubstraitTypeParser.ListContext):
-                type =  _evaluate(parametrized_type.expr(),values)
+                child_type = _evaluate(parametrized_type.expr(), values)
                 return Type(
                     list=Type.List(
-                        type=type,
+                        type=child_type,
                         nullability=nullability,
                     )
                 )
 
             elif isinstance(parametrized_type, SubstraitTypeParser.MapContext):
                 return Type(
                     map=Type.Map(
-                        key=_evaluate(parametrized_type.key,values),
-                        value=_evaluate(parametrized_type.value,values),
+                        key=_evaluate(parametrized_type.key, values),
+                        value=_evaluate(parametrized_type.value, values),
                         nullability=nullability,
                     )
                 )

src/substrait/extension_registry.py

@@ -1,17 +1,19 @@
-import yaml
 import itertools
 import re
-from substrait.gen.proto.type_pb2 import Type
-from importlib.resources import files as importlib_files
 from collections import defaultdict
+from importlib.resources import files as importlib_files
 from pathlib import Path
 from typing import Optional, Union
-from .derivation_expression import evaluate, _evaluate, _parse
+
+import yaml
+
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
 from substrait.gen.json import simple_extensions as se
+from substrait.gen.proto.type_pb2 import Type
 from substrait.simple_extension_utils import build_simple_extensions
-from .bimap import UriUrnBiDiMap
 
+from .bimap import UriUrnBiDiMap
+from .derivation_expression import _evaluate, _parse, evaluate
 
 DEFAULT_URN_PREFIX = "https://github.com/substrait-io/substrait/blob/main/extensions"
 
@@ -290,7 +292,7 @@ def __init__(
     def __repr__(self) -> str:
         return f"{self.name}:{'_'.join(self.normalized_inputs)}"
 
-    def satisfies_signature(self, signature: tuple) -> Optional[str]:
+    def satisfies_signature(self, signature: tuple | list) -> Optional[str]:
         if self.impl.variadic:
             min_args_allowed = self.impl.variadic.min or 0
             if len(signature) < min_args_allowed:
@@ -322,14 +324,12 @@ def satisfies_signature(self, signature: tuple) -> Optional[str]:
         output_type = evaluate(self.impl.return_, parameters)
 
         if self.nullability == se.NullabilityHandling.MIRROR:
-            sig_contains_nullable = any(
-                [
-                    p.__getattribute__(p.WhichOneof("kind")).nullability
-                    == Type.NULLABILITY_NULLABLE
-                    for p in signature
-                    if isinstance(p, Type)
-                ]
-            )
+            sig_contains_nullable = any([
+                p.__getattribute__(p.WhichOneof("kind")).nullability
+                == Type.NULLABILITY_NULLABLE
+                for p in signature
+                if isinstance(p, Type)
+            ])
             output_type.__getattribute__(output_type.WhichOneof("kind")).nullability = (
                 Type.NULLABILITY_NULLABLE
                 if sig_contains_nullable
@@ -417,7 +417,7 @@ def register_extension_dict(self, definitions: dict, uri: str) -> None:
 
     # TODO add an optional return type check
     def lookup_function(
-        self, urn: str, function_name: str, signature: tuple
+        self, urn: str, function_name: str, signature: tuple[Type] | list[Type]
     ) -> Optional[tuple[FunctionEntry, Type]]:
         if (
             urn not in self._function_mapping

tests/test_extension_registry.py

@@ -572,4 +572,4 @@ def test_register_requires_uri():
 
     # During migration, URI is required - this should fail with TypeError
     with pytest.raises(TypeError):
-        registry.register_extension_dict(yaml.safe_load(content))
+        registry.register_extension_dict(yaml.safe_load(content))