refactor: apply ruff fixes

vbarua · vbarua · commit ae44e6cd0b8b · 2026-01-13T16:40:46.000-08:00
diff --git a/examples/builder_example.py b/examples/builder_example.py
@@ -1,4 +1,5 @@
 import substrait.algebra_pb2 as stalg
+
 from substrait.builders.extended_expression import (
     column,
     literal,
diff --git a/src/substrait/builders/extended_expression.py b/src/substrait/builders/extended_expression.py
@@ -6,6 +6,7 @@
 import substrait.extended_expression_pb2 as stee
 import substrait.extensions.extensions_pb2 as ste
 import substrait.type_pb2 as stp
+
 from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_extended_expression_schema
 from substrait.utils import (
diff --git a/src/substrait/builders/plan.py b/src/substrait/builders/plan.py
@@ -12,19 +12,20 @@
 import substrait.extended_expression_pb2 as stee
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
-from substrait.version import __substrait_version__
+from substrait.extensions.extensions_pb2 import AdvancedExtension
+
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     resolve_expression,
 )
 from substrait.extension_registry import ExtensionRegistry
-from substrait.extensions.extensions_pb2 import AdvancedExtension
 from substrait.type_inference import infer_plan_schema
 from substrait.utils import (
     merge_extension_declarations,
     merge_extension_uris,
     merge_extension_urns,
 )
+from substrait.version import __substrait_version__
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 
diff --git a/src/substrait/dataframe/expression.py b/src/substrait/dataframe/expression.py
@@ -1,5 +1,6 @@
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stp
+
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     UnboundExtendedExpression,
diff --git a/src/substrait/derivation_expression.py b/src/substrait/derivation_expression.py
@@ -1,10 +1,10 @@
 from typing import Optional
 
 from antlr4 import CommonTokenStream, InputStream
+from substrait.type_pb2 import NamedStruct, Type
 
 from substrait.gen.antlr.SubstraitTypeLexer import SubstraitTypeLexer
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
-from substrait.type_pb2 import NamedStruct, Type
 
 
 def _evaluate(x, values: dict):
diff --git a/src/substrait/extension_registry/function_entry.py b/src/substrait/extension_registry/function_entry.py
@@ -3,9 +3,10 @@
 from enum import Enum
 from typing import Optional, Union
 
+from substrait.type_pb2 import Type
+
 from substrait.derivation_expression import _parse, evaluate
 from substrait.gen.json import simple_extensions as se
-from substrait.type_pb2 import Type
 
 from .signature_checker_helpers import covers, normalize_substrait_type_names
 
diff --git a/src/substrait/extension_registry/registry.py b/src/substrait/extension_registry/registry.py
@@ -8,9 +8,9 @@
 from typing import Optional, Union
 
 import yaml
+from substrait.type_pb2 import Type
 
 from substrait.bimap import UriUrnBiDiMap
-from substrait.type_pb2 import Type
 from substrait.simple_extension_utils import build_simple_extensions
 
 from .function_entry import FunctionEntry, FunctionType
diff --git a/src/substrait/extension_registry/signature_checker_helpers.py b/src/substrait/extension_registry/signature_checker_helpers.py
@@ -2,9 +2,10 @@
 
 from typing import Dict
 
+from substrait.type_pb2 import Type
+
 from substrait.derivation_expression import _evaluate
 from substrait.gen.antlr.SubstraitTypeParser import SubstraitTypeParser
-from substrait.type_pb2 import Type
 
 from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError
 
diff --git a/src/substrait/proto/__init__.py b/src/substrait/proto/__init__.py
@@ -5,15 +5,14 @@
 import substrait.function_pb2 as function
 import substrait.parameterized_types_pb2 as parameterized_types
 import substrait.plan_pb2 as plan
-import substrait.type_pb2 as type
 import substrait.type_expressions_pb2 as type_expressions
-
+import substrait.type_pb2 as type
 from substrait.algebra_pb2 import *
 from substrait.capabilities_pb2 import *
 from substrait.extended_expression_pb2 import *
 from substrait.extensions.extensions_pb2 import *
 from substrait.function_pb2 import *
 from substrait.parameterized_types_pb2 import *
 from substrait.plan_pb2 import *
+from substrait.type_expressions_pb2 import *
 from substrait.type_pb2 import *
-from substrait.type_expressions_pb2 import *
diff --git a/src/substrait/sql/sql_to_substrait.py b/src/substrait/sql/sql_to_substrait.py
@@ -5,6 +5,8 @@
 from deepdiff import DeepDiff
 from sqloxide import parse_sql
 
+from substrait import algebra_pb2 as stalg
+from substrait import type_pb2 as stt
 from substrait.builders.extended_expression import (
     UnboundExtendedExpression,
     aggregate_function,
@@ -24,8 +26,6 @@
     sort,
 )
 from substrait.extension_registry import ExtensionRegistry
-from substrait import algebra_pb2 as stalg
-from substrait import type_pb2 as stt
 
 SchemaResolver = Callable[[str], stt.NamedStruct]
 
diff --git a/tests/builders/extended_expression/test_aggregate_function.py b/tests/builders/extended_expression/test_aggregate_function.py
@@ -1,9 +1,9 @@
-import yaml
-
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.extensions.extensions_pb2 as ste
 import substrait.type_pb2 as stt
+import yaml
+
 from substrait.builders.extended_expression import aggregate_function, literal
 from substrait.extension_registry import ExtensionRegistry
 
diff --git a/tests/builders/extended_expression/test_cast.py b/tests/builders/extended_expression/test_cast.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import cast, literal
 from substrait.builders.type import i8, i16
 from substrait.extension_registry import ExtensionRegistry
@@ -47,9 +48,9 @@ def test_cast():
 
 
 def test_cast_with_extension():
+    import substrait.extensions.extensions_pb2 as ste
     import yaml
 
-    import substrait.extensions.extensions_pb2 as ste
     from substrait.builders.extended_expression import scalar_function
 
     registry_with_ext = ExtensionRegistry(load_default_extensions=False)
diff --git a/tests/builders/extended_expression/test_column.py b/tests/builders/extended_expression/test_column.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import column
 
 struct = stt.Type.Struct(
diff --git a/tests/builders/extended_expression/test_if_then.py b/tests/builders/extended_expression/test_if_then.py
@@ -2,6 +2,7 @@
 import substrait.extended_expression_pb2 as stee
 import substrait.extensions.extensions_pb2 as ste
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import (
     column,
     if_then,
diff --git a/tests/builders/extended_expression/test_literal.py b/tests/builders/extended_expression/test_literal.py
@@ -1,6 +1,7 @@
 from datetime import date
 
 import substrait.algebra_pb2 as stalg
+
 from substrait.builders import type as sttb
 from substrait.builders.extended_expression import literal
 
diff --git a/tests/builders/extended_expression/test_multi_or_list.py b/tests/builders/extended_expression/test_multi_or_list.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal, multi_or_list
 from substrait.builders.type import i8
 from substrait.extension_registry import ExtensionRegistry
@@ -84,9 +85,9 @@ def test_singular_or_list():
 
 
 def test_multi_or_list_with_extension():
+    import substrait.extensions.extensions_pb2 as ste
     import yaml
 
-    import substrait.extensions.extensions_pb2 as ste
     from substrait.builders.extended_expression import scalar_function
 
     registry_with_ext = ExtensionRegistry(load_default_extensions=False)
diff --git a/tests/builders/extended_expression/test_scalar_function.py b/tests/builders/extended_expression/test_scalar_function.py
@@ -1,9 +1,9 @@
-import yaml
-
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.extensions.extensions_pb2 as ste
 import substrait.type_pb2 as stt
+import yaml
+
 from substrait.builders.extended_expression import literal, scalar_function
 from substrait.extension_registry import ExtensionRegistry
 
diff --git a/tests/builders/extended_expression/test_singular_or_list.py b/tests/builders/extended_expression/test_singular_or_list.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal, singular_or_list
 from substrait.builders.type import i8
 from substrait.extension_registry import ExtensionRegistry
@@ -53,9 +54,9 @@ def test_singular_or_list():
 
 
 def test_singular_or_list_with_extension():
+    import substrait.extensions.extensions_pb2 as ste
     import yaml
 
-    import substrait.extensions.extensions_pb2 as ste
     from substrait.builders.extended_expression import scalar_function
 
     registry_with_ext = ExtensionRegistry(load_default_extensions=False)
diff --git a/tests/builders/extended_expression/test_switch.py b/tests/builders/extended_expression/test_switch.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal, switch
 from substrait.builders.type import i8
 from substrait.extension_registry import ExtensionRegistry
@@ -77,9 +78,9 @@ def test_switch():
 
 
 def test_switch_with_extension():
+    import substrait.extensions.extensions_pb2 as ste
     import yaml
 
-    import substrait.extensions.extensions_pb2 as ste
     from substrait.builders.extended_expression import scalar_function
 
     registry_with_ext = ExtensionRegistry(load_default_extensions=False)
diff --git a/tests/builders/extended_expression/test_window_function.py b/tests/builders/extended_expression/test_window_function.py
@@ -1,9 +1,9 @@
-import yaml
-
 import substrait.algebra_pb2 as stalg
 import substrait.extended_expression_pb2 as stee
 import substrait.extensions.extensions_pb2 as ste
 import substrait.type_pb2 as stt
+import yaml
+
 from substrait.builders.extended_expression import window_function
 from substrait.extension_registry import ExtensionRegistry
 
diff --git a/tests/builders/plan/test_aggregate.py b/tests/builders/plan/test_aggregate.py
@@ -1,9 +1,9 @@
-import yaml
-
 import substrait.algebra_pb2 as stalg
 import substrait.extensions.extensions_pb2 as ste
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+import yaml
+
 from substrait.builders.extended_expression import aggregate_function, column
 from substrait.builders.plan import aggregate, default_version, read_named_table
 from substrait.builders.type import boolean, i64
diff --git a/tests/builders/plan/test_cross.py b/tests/builders/plan/test_cross.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.plan import cross, default_version, read_named_table
 from substrait.builders.type import boolean, i64, string
 from substrait.extension_registry import ExtensionRegistry
diff --git a/tests/builders/plan/test_fetch.py b/tests/builders/plan/test_fetch.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal
 from substrait.builders.plan import default_version, fetch, read_named_table
 from substrait.builders.type import boolean, i64
diff --git a/tests/builders/plan/test_filter.py b/tests/builders/plan/test_filter.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal
 from substrait.builders.plan import default_version, filter, read_named_table
 from substrait.builders.type import boolean, i64
diff --git a/tests/builders/plan/test_join.py b/tests/builders/plan/test_join.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import literal
 from substrait.builders.plan import default_version, join, read_named_table
 from substrait.builders.type import boolean, i64, string
diff --git a/tests/builders/plan/test_project.py b/tests/builders/plan/test_project.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import column
 from substrait.builders.plan import default_version, project, read_named_table, select
 from substrait.builders.type import boolean, i64
diff --git a/tests/builders/plan/test_read.py b/tests/builders/plan/test_read.py
@@ -1,13 +1,13 @@
 import pytest
-from google.protobuf import any_pb2
-from google.protobuf.wrappers_pb2 import StringValue
-
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+from google.protobuf import any_pb2
+from google.protobuf.wrappers_pb2 import StringValue
+from substrait.extensions.extensions_pb2 import AdvancedExtension
+
 from substrait.builders.plan import default_version, read_named_table
 from substrait.builders.type import boolean, i64
-from substrait.extensions.extensions_pb2 import AdvancedExtension
 
 struct = stt.Type.Struct(
     types=[i64(nullable=False), boolean()],
diff --git a/tests/builders/plan/test_set.py b/tests/builders/plan/test_set.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.plan import default_version, read_named_table, set
 from substrait.builders.type import boolean, i64
 from substrait.extension_registry import ExtensionRegistry
diff --git a/tests/builders/plan/test_sort.py b/tests/builders/plan/test_sort.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.extended_expression import column
 from substrait.builders.plan import default_version, read_named_table, sort
 from substrait.builders.type import boolean, i64
diff --git a/tests/builders/plan/test_write.py b/tests/builders/plan/test_write.py
@@ -1,6 +1,7 @@
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
 from substrait.builders.plan import read_named_table, write_named_table
 from substrait.builders.type import boolean, i64
 
diff --git a/tests/builders/test_types.py b/tests/builders/test_types.py
@@ -1,6 +1,6 @@
 import pytest
-
 import substrait.type_pb2 as stt
+
 from substrait.builders.type import boolean, i64, named_struct
 
 
diff --git a/tests/dataframe/test_df_project.py b/tests/dataframe/test_df_project.py
@@ -1,7 +1,8 @@
-import substrait.dataframe as sdf
 import substrait.algebra_pb2 as stalg
 import substrait.plan_pb2 as stp
 import substrait.type_pb2 as stt
+
+import substrait.dataframe as sdf
 from substrait.builders.plan import default_version, read_named_table
 from substrait.builders.type import boolean, i64
 from substrait.extension_registry import ExtensionRegistry
diff --git a/tests/extension_registry/test_registry_lookup.py b/tests/extension_registry/test_registry_lookup.py
@@ -1,6 +1,7 @@
-from substrait.builders.type import decimal, i8, i16
 from substrait.type_pb2 import Type
 
+from substrait.builders.type import decimal, i8, i16
+
 
 def test_non_existing_urn(registry):
     assert (
diff --git a/tests/extension_registry/test_type_coverage.py b/tests/extension_registry/test_type_coverage.py
@@ -1,11 +1,12 @@
 """Tests for the covers() function - type coverage and matching."""
 
+from substrait.type_pb2 import Type
+
 from substrait.builders.type import decimal, i8, i16, i32, struct
 from substrait.builders.type import list as list_
 from substrait.builders.type import map as map_
 from substrait.derivation_expression import _parse
 from substrait.extension_registry import covers
-from substrait.type_pb2 import Type
 
 
 def test_covers():
diff --git a/tests/test_derivation_expression.py b/tests/test_derivation_expression.py
diff --git a/tests/test_literal_type_inference.py b/tests/test_literal_type_inference.py
diff --git a/tests/test_type_inference.py b/tests/test_type_inference.py
diff --git a/tests/test_uri_urn_migration.py b/tests/test_uri_urn_migration.py
diff --git a/tests/test_utils.py b/tests/test_utils.py

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,5 @@`
`1`	`1`	`import substrait.algebra_pb2 as stalg`
	`2`	`+`
`2`	`3`	`from substrait.builders.extended_expression import (`
`3`	`4`	`column,`
`4`	`5`	`literal,`