Repository: mrpowers-io/quinn (language: python)

Issue:
Add a method to print the DataFrame schema in code format
https://github.com/MrPowers/spark-daria/blob/master/src/main/scala/com/github/mrpowers/spark/daria/sql/DataFrameExt.scala


Discussion / hints:

Fix patch:
diff --git a/quinn/__init__.py b/quinn/__init__.py
index 2008de62..5aa62b38 100644
--- a/quinn/__init__.py
+++ b/quinn/__init__.py
@@ -1,5 +1,6 @@
 from .dataframe_helpers import *
 from .dataframe_validator import *
+from .schema_helpers import print_schema_as_code
 from .functions import *
 from .scala_to_pyspark import ScalaToPyspark
 from .transformations import *
diff --git a/quinn/schema_helpers.py b/quinn/schema_helpers.py
new file mode 100644
index 00000000..61016bed
--- /dev/null
+++ b/quinn/schema_helpers.py
@@ -0,0 +1,72 @@
+from pyspark.sql import types as T
+
+
+def print_schema_as_code(dtype: T.DataType) -> str:
+    """Represent DataType (including StructType) as valid Python code.
+
+    Parameters
+    ----------
+    dtype: pyspark.sql.types.DataType
+        The input DataType or Schema object
+
+    Returns
+    -------
+    str
+        A valid python code which generate the same schema.
+    """
+    res = []
+    if isinstance(dtype, T.StructType):
+        res.append("StructType(\n\tfields=[")
+        for field in dtype.fields:
+            for line in _repr_column(field).split("\n"):
+                res.append("\n\t\t")
+                res.append(line)
+            res.append(",")
+        res.append("\n\t]\n)")
+
+    elif isinstance(dtype, T.ArrayType):
+        res.append("ArrayType(")
+        res.append(print_schema_as_code(dtype.elementType))
+        res.append(")")
+
+    elif isinstance(dtype, T.MapType):
+        res.append("MapType(")
+        res.append(f"\n\t{print_schema_as_code(dtype.keyType)},")
+        for line in print_schema_as_code(dtype.valueType).split("\n"):
+            res.append("\n\t")
+            res.append(line)
+        res.append(",")
+        res.append(f"\n\t{dtype.valueContainsNull},")
+        res.append("\n)")
+
+    elif isinstance(dtype, T.DecimalType):
+        res.append(f"DecimalType({dtype.precision}, {dtype.scale})")
+
+    else:
+        res.append(f"{dtype}()")
+
+    return "".join(res)
+
+
+def _repr_column(column: T.StructField) -> str:
+    res = []
+
+    if (
+        isinstance(column.dataType, T.StructType)
+        or isinstance(column.dataType, T.ArrayType)
+        or isinstance(column.dataType, T.MapType)
+    ):
+        res.append(f'StructField(\n\t"{column.name}",')
+        for line in print_schema_as_code(column.dataType).split("\n"):
+            res.append("\n\t")
+            res.append(line)
+        res.append(",")
+        res.append(f"\n\t{column.nullable},")
+        res.append("\n)")
+
+    else:
+        res.append(
+            f'StructField("{column.name}", {print_schema_as_code(column.dataType)}, {column.nullable})'
+        )
+
+    return "".join(res)


Test patch:
diff --git a/tests/extensions/test_dataframe_ext.py b/tests/extensions/test_dataframe_ext.py
index 19b0e412..e853d6c9 100644
--- a/tests/extensions/test_dataframe_ext.py
+++ b/tests/extensions/test_dataframe_ext.py
@@ -1,11 +1,13 @@
 from functools import partial
 
+import chispa
+from chispa.schema_comparer import assert_basic_schema_equality
+from pyspark.sql.dataframe import DataFrame
 from pyspark.sql.functions import col
 
 from tests.conftest import auto_inject_fixtures
-from .dataframe_transformations import *
 
-import chispa
+from .dataframe_transformations import *
 
 
 @auto_inject_fixtures("spark")
diff --git a/tests/test_schema_helpers.py b/tests/test_schema_helpers.py
new file mode 100644
index 00000000..56959199
--- /dev/null
+++ b/tests/test_schema_helpers.py
@@ -0,0 +1,31 @@
+from pyspark.sql.types import *
+
+from quinn.schema_helpers import print_schema_as_code
+
+from chispa.schema_comparer import assert_basic_schema_equality
+
+from tests.conftest import auto_inject_fixtures
+
+
+@auto_inject_fixtures("spark")
+def test_print_schema_as_code(spark):
+    fields = []
+    fields.append(StructField("simple_int", IntegerType()))
+    fields.append(StructField("decimal_with_nums", DecimalType(19, 8)))
+    fields.append(StructField("array", ArrayType(FloatType())))
+    fields.append(StructField("map", MapType(StringType(), ArrayType(DoubleType()))))
+    fields.append(
+        StructField(
+            "struct",
+            StructType(
+                [
+                    StructField("first", StringType()),
+                    StructField("second", TimestampType()),
+                ]
+            ),
+        )
+    )
+
+    schema = StructType(fields=fields)
+
+    assert_basic_schema_equality(schema, eval(print_schema_as_code(schema)))