Moved logic to convert Option class to dict

astronomer · Jan 6, 2023 · f570616 · f570616
1 parent 11fafb2
commit f570616
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 8 deletions.
diff --git a/python-sdk/src/astro/files/types/csv.py b/python-sdk/src/astro/files/types/csv.py
@@ -2,7 +2,6 @@
 
 import io
 
-import attr
 import pandas as pd
 
 from astro.constants import FileType as FileTypeConstants
@@ -32,7 +31,7 @@ def export_to_dataframe(
             in the resulting dataframe
         """
         if isinstance(load_options, PandasLoadOptions):
-            kwargs.update(attr.asdict(load_options))
+            kwargs.update(load_options.to_dict())
         df = pd.read_csv(stream, **kwargs)
         df = convert_columns_names_capitalization(
             df=df, columns_names_capitalization=columns_names_capitalization

diff --git a/python-sdk/src/astro/files/types/json.py b/python-sdk/src/astro/files/types/json.py
@@ -2,7 +2,6 @@
 
 import io
 
-import attr
 import pandas as pd
 
 from astro.constants import FileType as FileTypeConstants
@@ -35,7 +34,7 @@ def export_to_dataframe(
         # Pandas `read_json` does not support the `nrows` parameter unless we're using NDJSON
         kwargs_copy.pop("nrows", None)
         if isinstance(load_options, PandasLoadOptions):
-            kwargs_copy.update(attr.asdict(load_options))
+            kwargs_copy.update(load_options.to_dict())
         df = pd.read_json(stream, **kwargs_copy)
         df = convert_columns_names_capitalization(
             df=df, columns_names_capitalization=columns_names_capitalization

diff --git a/python-sdk/src/astro/files/types/ndjson.py b/python-sdk/src/astro/files/types/ndjson.py
@@ -3,7 +3,6 @@
 import io
 import json
 
-import attr
 import pandas as pd
 
 from astro.constants import DEFAULT_CHUNK_SIZE, FileType as FileTypeConstants
@@ -32,7 +31,7 @@ def export_to_dataframe(
             in the resulting dataframe
         """
         if isinstance(load_options, PandasLoadOptions):
-            kwargs.update(attr.asdict(load_options))
+            kwargs.update(load_options.to_dict())
         df = NDJSONFileType.flatten(self.normalize_config, stream, **kwargs)
         df = convert_columns_names_capitalization(
             df=df, columns_names_capitalization=columns_names_capitalization

diff --git a/python-sdk/src/astro/files/types/parquet.py b/python-sdk/src/astro/files/types/parquet.py
@@ -2,7 +2,6 @@
 
 import io
 
-import attr
 import pandas as pd
 
 from astro.constants import FileType as FileTypeConstants
@@ -36,7 +35,7 @@ def export_to_dataframe(
 
         byte_io_buffer = self._convert_remote_file_to_byte_stream(stream)
         if isinstance(load_options, PandasLoadOptions):
-            kwargs_copy.update(attr.asdict(load_options))
+            kwargs_copy.update(load_options.to_dict())
 
         df = pd.read_parquet(byte_io_buffer, **kwargs_copy)
         df = convert_columns_names_capitalization(

diff --git a/python-sdk/src/astro/options.py b/python-sdk/src/astro/options.py
@@ -5,3 +5,9 @@
 class LoadOptions:
     def empty(self):
         return NotImplementedError()
+
+    def to_dict(self) -> dict:
+        """
+        Convert options class to dict
+        """
+        return attr.asdict(self)