fix: ruff and type hints

Author: giospada

src/substrait/builders/plan.py

@@ -5,29 +5,38 @@
 See `examples/builder_example.py` for usage.
 """
 
+from typing import Callable, Iterable, Optional, TypedDict, Union
+
+import substrait.gen.proto.algebra_pb2 as stalg
 import substrait.gen.proto.extended_expression_pb2 as stee
-from typing import Iterable, Optional, Union, Callable
-from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
 import substrait.gen.proto.plan_pb2 as stp
 import substrait.gen.proto.type_pb2 as stt
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     resolve_expression,
 )
 from substrait.extension_registry import ExtensionRegistry
+from substrait.gen.proto.extensions.extensions_pb2 import AdvancedExtension
 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

@@ -8,11 +8,6 @@
     i8,
     i16,
     decimal,
-    var_char,
-    fixed_char,
-    fixed_binary,
-    timestamp,
-    precision_timestamp_tz,
 )
 
 content = """%YAML 1.2
@@ -193,9 +188,9 @@ def test_wildcard_match_with_constraits():
         registry.lookup_function(
             urn="extension:test:functions",
             function_name="add",
-            signature=[i16((nullable=False), i16((nullable=False), i8((nullable=False)],
+            signature=[i16(nullable=False), i16(nullable=False), i8(nullable=False)],
         )[1]
-        == i8((nullable=False)
+        == i8(nullable=False)
     )
 
 
@@ -204,9 +199,9 @@ def test_variadic():
         registry.lookup_function(
             urn="extension:test:functions",
             function_name="test_fn",
-            signature=[i8((nullable=False), i8((nullable=False), i8((nullable=False)],
+            signature=[i8(nullable=False), i8(nullable=False), i8(nullable=False)],
         )[1]
-        == i8((nullable=False)
+        == i8(nullable=False)
     )
 
 
@@ -215,9 +210,9 @@ def test_variadic_any():
         registry.lookup_function(
             urn="extension:test:functions",
             function_name="test_fn_variadic_any",
-            signature=[i16((nullable=False), i16((nullable=False), i16((nullable=False)],
+            signature=[i16(nullable=False), i16(nullable=False), i16(nullable=False)],
         )[1]
-        == i16((nullable=False)
+        == i16(nullable=False)
     )
 
 
@@ -577,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))