fix: formatting and linter errors

Author: giospada

src/substrait/builders/type.py

@@ -221,6 +221,7 @@ def precision_timestamp_tz(precision: int, nullable=True) -> stt.Type:
         )
     )
 
+
 def timestamp(nullable=True) -> stt.Type:
     return stt.Type(
         timestamp=stt.Type.Timestamp(
@@ -230,6 +231,7 @@ def timestamp(nullable=True) -> stt.Type:
         )
     )
 
+
 def struct(types: Iterable[stt.Type], nullable=True) -> stt.Type:
     return stt.Type(
         struct=stt.Type.Struct(

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
@@ -121,15 +123,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 +150,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)
+                list_type = _evaluate(parametrized_type.expr(), values)
                 return Type(
                     list=Type.List(
-                        type=type,
+                        type=list_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,
                     )
                 )

tests/test_extension_registry.py

@@ -129,11 +129,12 @@
 )
 
 
-
 def test_non_existing_urn():
     assert (
         registry.lookup_function(
-            urn="non_existent", function_name="add", signature=[i8(nullable=False), i8(nullable=False)]
+            urn="non_existent",
+            function_name="add",
+            signature=[i8(nullable=False), i8(nullable=False)],
         )
         is None
     )
@@ -142,8 +143,9 @@ def test_non_existing_urn():
 def test_non_existing_function():
     assert (
         registry.lookup_function(
-
-        urn="extension:test:functions", function_name="sub", signature=[i8(nullable=False), i8(nullable=False)]
+            urn="extension:test:functions",
+            function_name="sub",
+            signature=[i8(nullable=False), i8(nullable=False)],
         )
         is None
     )
@@ -152,15 +154,19 @@ def test_non_existing_function():
 def test_non_existing_function_signature():
     assert (
         registry.lookup_function(
-            urn="extension:test:functions", function_name="add", signature=[i8(nullable=False)]
+            urn="extension:test:functions",
+            function_name="add",
+            signature=[i8(nullable=False)],
         )
         is None
     )
 
 
 def test_exact_match():
     assert registry.lookup_function(
-        urn="extension:test:functions", function_name="add", signature=[i8(nullable=False), i8(nullable=False)]
+        urn="extension:test:functions",
+        function_name="add",
+        signature=[i8(nullable=False), i8(nullable=False)],
     )[1] == Type(i8=Type.I8(nullability=Type.NULLABILITY_REQUIRED))
 
 
@@ -184,43 +190,35 @@ def test_wildcard_match_fails_with_constraits():
 
 
 def test_wildcard_match_with_constraits():
-    assert (
-        registry.lookup_function(
-            urn="extension:test:functions",
-            function_name="add",
-            signature=[i16(nullable=False), i16(nullable=False), i8(nullable=False)],
-        )[1]
-        == i8(nullable=False)
-    )
+    assert registry.lookup_function(
+        urn="extension:test:functions",
+        function_name="add",
+        signature=[i16(nullable=False), i16(nullable=False), i8(nullable=False)],
+    )[1] == i8(nullable=False)
 
 
 def test_variadic():
-    assert (
-        registry.lookup_function(
-            urn="extension:test:functions",
-            function_name="test_fn",
-            signature=[i8(nullable=False), i8(nullable=False), i8(nullable=False)],
-        )[1]
-        == i8(nullable=False)
-    )
+    assert registry.lookup_function(
+        urn="extension:test:functions",
+        function_name="test_fn",
+        signature=[i8(nullable=False), i8(nullable=False), i8(nullable=False)],
+    )[1] == i8(nullable=False)
 
 
 def test_variadic_any():
-    assert (
-        registry.lookup_function(
-            urn="extension:test:functions",
-            function_name="test_fn_variadic_any",
-            signature=[i16(nullable=False), i16(nullable=False), i16(nullable=False)],
-        )[1]
-        == i16(nullable=False)
-    )
+    assert registry.lookup_function(
+        urn="extension:test:functions",
+        function_name="test_fn_variadic_any",
+        signature=[i16(nullable=False), i16(nullable=False), i16(nullable=False)],
+    )[1] == i16(nullable=False)
 
 
 def test_variadic_fails_min_constraint():
     assert (
         registry.lookup_function(
-
-            urn="extension:test:functions", function_name="test_fn", signature=[i8(nullable=False)]
+            urn="extension:test:functions",
+            function_name="test_fn",
+            signature=[i8(nullable=False)],
         )
         is None
     )
@@ -254,14 +252,11 @@ def test_decimal_happy_path_discrete():
 
 
 def test_enum_with_valid_option():
-    assert (
-        registry.lookup_function(
-            urn="extension:test:functions",
-            function_name="test_enum",
-            signature=["FLIP", i8(nullable=False)],
-        )[1]
-        == i8(nullable=False)
-    )
+    assert registry.lookup_function(
+        urn="extension:test:functions",
+        function_name="test_enum",
+        signature=["FLIP", i8(nullable=False)],
+    )[1] == i8(nullable=False)
 
 
 def test_enum_with_nonexistent_option():