format

Author: chenkovsky

python/datafusion/__init__.py

@@ -26,7 +26,7 @@
 except ImportError:
     import importlib_metadata
 
-from . import functions, object_store, substrait
+from . import functions, object_store, substrait, unparser
 
 # The following imports are okay to remain as opaque to the user.
 from ._internal import Config
@@ -49,7 +49,6 @@
 from .plan import ExecutionPlan, LogicalPlan
 from .record_batch import RecordBatch, RecordBatchStream
 from .udf import Accumulator, AggregateUDF, ScalarUDF, WindowUDF, udaf, udf, udwf
-from .unparser import Dialect, Unparser
 
 __version__ = importlib_metadata.version(__name__)
 
@@ -61,7 +60,6 @@
     "DFSchema",
     "DataFrame",
     "Database",
-    "Dialect",
     "ExecutionPlan",
     "Expr",
     "LogicalPlan",
@@ -73,7 +71,6 @@
     "SessionConfig",
     "SessionContext",
     "Table",
-    "Unparser",
     "WindowFrame",
     "WindowUDF",
     "col",
@@ -92,6 +89,7 @@
     "udaf",
     "udf",
     "udwf",
+    "unparser",
 ]
 
 

python/datafusion/unparser.py

@@ -31,6 +31,11 @@ def __init__(self, dialect: unparser_internal.Dialect) -> None:
         """This constructor is not typically called by the end user."""
         self.dialect = dialect
 
+    @staticmethod
+    def default() -> "Dialect":
+        """Create a new default dialect."""
+        return Dialect(unparser_internal.Dialect.default())
+
     @staticmethod
     def mysql() -> "Dialect":
         """Create a new MySQL dialect."""
@@ -63,6 +68,11 @@ def plan_to_sql(self, plan: LogicalPlan) -> str:
         """Convert a logical plan to a SQL string."""
         return self.unparser.plan_to_sql(plan._raw_plan)
 
+    def with_pretty(self, pretty: bool) -> "Unparser":
+        """Set the pretty flag."""
+        self.unparser = self.unparser.with_pretty(pretty)
+        return self
+
 
 __all__ = [
     "Dialect",