fix: resolve review comments

eakmanrq · Nov 4, 2024 · ad64f1e · ad64f1e
1 parent f27c652
commit ad64f1e
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/sqlframe/base/mixins/readwriter_mixins.py b/sqlframe/base/mixins/readwriter_mixins.py
@@ -75,8 +75,7 @@ def load(
 
         assert path is not None, "path is required"
         assert isinstance(path, str), "path must be a string"
-        format = format or self.state_format_to_read
-        format = format or _infer_format(path)
+        format = format or self.state_format_to_read or _infer_format(path)
         kwargs = {k: v for k, v in options.items() if v is not None}
         if format == "json":
             df = pd.read_json(path, lines=True, **kwargs)  # type: ignore

diff --git a/sqlframe/base/readerwriter.py b/sqlframe/base/readerwriter.py
@@ -36,7 +36,7 @@
 class _BaseDataFrameReader(t.Generic[SESSION, DF]):
     def __init__(self, spark: SESSION):
         self._session = spark
-        self.state_format_to_read = None
+        self.state_format_to_read: t.Optional[str] = None
 
     @property
     def session(self) -> SESSION: