refactor: update substrait.gen.proto imports to just substrait

Author: vbarua

examples/builder_example.py

@@ -1,4 +1,5 @@
-import substrait.gen.proto.algebra_pb2 as stalg
+import substrait.algebra_pb2 as stalg
+
 from substrait.builders.extended_expression import (
     column,
     literal,

src/substrait/builders/extended_expression.py

@@ -2,10 +2,11 @@
 from datetime import date
 from typing import Any, Callable, Iterable, Union
 
-import substrait.gen.proto.algebra_pb2 as stalg
-import substrait.gen.proto.extended_expression_pb2 as stee
-import substrait.gen.proto.extensions.extensions_pb2 as ste
-import substrait.gen.proto.type_pb2 as stp
+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 stp
+
 from substrait.extension_registry import ExtensionRegistry
 from substrait.type_inference import infer_extended_expression_schema
 from substrait.utils import (

src/substrait/builders/plan.py

@@ -8,23 +8,24 @@
 import re
 from typing import Callable, Iterable, Optional, Union
 
-import substrait.gen.proto.algebra_pb2 as stalg
-import substrait.gen.proto.extended_expression_pb2 as stee
-import substrait.gen.proto.plan_pb2 as stp
-import substrait.gen.proto.type_pb2 as stt
-from substrait.version import __substrait_version__
+import substrait.algebra_pb2 as stalg
+import substrait.extended_expression_pb2 as stee
+import substrait.plan_pb2 as stp
+import substrait.type_pb2 as stt
+from substrait.extensions.extensions_pb2 import AdvancedExtension
+
 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_uris,
     merge_extension_urns,
 )
+from substrait.version import __substrait_version__
 
 UnboundPlan = Callable[[ExtensionRegistry], stp.Plan]
 

src/substrait/builders/type.py

@@ -1,6 +1,6 @@
 from typing import Iterable
 
-import substrait.gen.proto.type_pb2 as stt
+import substrait.type_pb2 as stt
 
 
 def boolean(nullable=True) -> stt.Type:

src/substrait/dataframe/expression.py

@@ -1,5 +1,6 @@
-import substrait.gen.proto.extended_expression_pb2 as stee
-import substrait.gen.proto.type_pb2 as stp
+import substrait.extended_expression_pb2 as stee
+import substrait.type_pb2 as stp
+
 from substrait.builders.extended_expression import (
     ExtendedExpressionOrUnbound,
     UnboundExtendedExpression,

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.gen.proto.type_pb2 import NamedStruct, Type
 
 
 def _evaluate(x, values: dict):

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.gen.proto.type_pb2 import Type
 
 from .signature_checker_helpers import covers, normalize_substrait_type_names
 

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.gen.proto.type_pb2 import Type
 from substrait.simple_extension_utils import build_simple_extensions
 
 from .function_entry import FunctionEntry, FunctionType

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.gen.proto.type_pb2 import Type
 
 from .exceptions import UnhandledParameterizedTypeError, UnrecognizedSubstraitTypeError
 

src/substrait/proto.py

@@ -8,7 +8,7 @@ def _load():
     the protocol modules with a friendly name making the protocol
     more convenient to use.
 
-    substrait.gen.proto.extensions.extensions_pb2.SimpleExtensionDeclaration
+    substrait.extensions.extensions_pb2.SimpleExtensionDeclaration
     becomes substrait.proto.SimpleExtensionDeclaration
     """
     import importlib