Expose the correct namespaces for Python modules

Author: datapythonista

python/datafusion/__init__.py

@@ -88,66 +88,6 @@
 
 __version__ = importlib_metadata.version(__name__)
 
-__all__ = [
-    "Config",
-    "DataFrame",
-    "SessionContext",
-    "SessionConfig",
-    "SQLOptions",
-    "RuntimeConfig",
-    "Expr",
-    "AggregateUDF",
-    "ScalarUDF",
-    "Window",
-    "WindowFrame",
-    "column",
-    "literal",
-    "TableScan",
-    "Projection",
-    "DFSchema",
-    "DFField",
-    "Analyze",
-    "Sort",
-    "Limit",
-    "Filter",
-    "Like",
-    "ILike",
-    "SimilarTo",
-    "ScalarVariable",
-    "Alias",
-    "Not",
-    "IsNotNull",
-    "IsTrue",
-    "IsFalse",
-    "IsUnknown",
-    "IsNotTrue",
-    "IsNotFalse",
-    "IsNotUnknown",
-    "Negative",
-    "ScalarFunction",
-    "BuiltinScalarFunction",
-    "InList",
-    "Exists",
-    "Subquery",
-    "InSubquery",
-    "ScalarSubquery",
-    "GroupingSet",
-    "Placeholder",
-    "Case",
-    "Cast",
-    "TryCast",
-    "Between",
-    "Explain",
-    "SubqueryAlias",
-    "Extension",
-    "CreateMemoryTable",
-    "CreateView",
-    "Distinct",
-    "DropTable",
-    "Repartition",
-    "Partitioning",
-]
-
 
 class Accumulator(metaclass=ABCMeta):
     @abstractmethod
@@ -175,6 +115,8 @@ def column(value):
 
 
 def literal(value):
+    import pyarrow as pa
+
     if not isinstance(value, pa.Scalar):
         value = pa.scalar(value)
     return Expr.literal(value)
@@ -204,6 +146,8 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None):
     """
     Create a new User Defined Aggregate Function
     """
+    import pyarrow as pa
+
     if not issubclass(accum, Accumulator):
         raise TypeError("`accum` must implement the abstract base class Accumulator")
     if name is None:
@@ -218,3 +162,68 @@ def udaf(accum, input_type, return_type, state_type, volatility, name=None):
         state_type=state_type,
         volatility=volatility,
     )
+
+
+del ABCMeta
+del abstractmethod
+del List
+del importlib_metadata
+del pa
+
+
+__all__ = [
+    "Config",
+    "DataFrame",
+    "SessionContext",
+    "SessionConfig",
+    "SQLOptions",
+    "RuntimeConfig",
+    "Expr",
+    "AggregateUDF",
+    "ScalarUDF",
+    "Window",
+    "WindowFrame",
+    "column",
+    "literal",
+    "TableScan",
+    "Projection",
+    "DFSchema",
+    "Analyze",
+    "Sort",
+    "Limit",
+    "Filter",
+    "Like",
+    "ILike",
+    "SimilarTo",
+    "ScalarVariable",
+    "Alias",
+    "Not",
+    "IsNotNull",
+    "IsTrue",
+    "IsFalse",
+    "IsUnknown",
+    "IsNotTrue",
+    "IsNotFalse",
+    "IsNotUnknown",
+    "Negative",
+    "InList",
+    "Exists",
+    "Subquery",
+    "InSubquery",
+    "ScalarSubquery",
+    "GroupingSet",
+    "Placeholder",
+    "Case",
+    "Cast",
+    "TryCast",
+    "Between",
+    "Explain",
+    "SubqueryAlias",
+    "Extension",
+    "CreateMemoryTable",
+    "CreateView",
+    "Distinct",
+    "DropTable",
+    "Repartition",
+    "Partitioning",
+]

python/datafusion/common.py

@@ -16,8 +16,15 @@
 # under the License.
 
 
-from ._internal import common
-
-
 def __getattr__(name):
+    from ._internal import common
+
     return getattr(common, name)
+
+
+def __dir__():
+    from ._internal import common
+
+    return list(globals().keys()) + [
+        obj for obj in dir(common) if not obj.startswith("_")
+    ]

python/datafusion/expr.py

@@ -16,8 +16,15 @@
 # under the License.
 
 
-from ._internal import expr
-
-
 def __getattr__(name):
+    from ._internal import expr
+
     return getattr(expr, name)
+
+
+def __dir__():
+    from ._internal import expr
+
+    return list(globals().keys()) + [
+        obj for obj in dir(expr) if not obj.startswith("_")
+    ]

python/datafusion/functions.py

@@ -16,8 +16,15 @@
 # under the License.
 
 
-from ._internal import functions
-
-
 def __getattr__(name):
+    from ._internal import functions
+
     return getattr(functions, name)
+
+
+def __dir__():
+    from ._internal import functions
+
+    return list(globals().keys()) + [
+        obj for obj in dir(functions) if not obj.startswith("_")
+    ]

python/datafusion/object_store.py

@@ -16,8 +16,15 @@
 # under the License.
 
 
-from ._internal import object_store
-
-
 def __getattr__(name):
+    from ._internal import object_store
+
     return getattr(object_store, name)
+
+
+def __dir__():
+    from ._internal import object_store
+
+    return list(globals().keys()) + [
+        obj for obj in dir(object_store) if not obj.startswith("_")
+    ]

python/datafusion/substrait.py

@@ -16,8 +16,15 @@
 # under the License.
 
 
-from ._internal import substrait
-
-
 def __getattr__(name):
+    from ._internal import substrait
+
     return getattr(substrait, name)
+
+
+def __dir__():
+    from ._internal import substrait
+
+    return list(globals().keys()) + [
+        obj for obj in dir(substrait) if not obj.startswith("_")
+    ]