keep type checker happy

Author: nickzoic

countess/plugins/csv.py

@@ -54,6 +54,19 @@ class ColumnsMultiParam(MultiParam):
 CSV_DELIMITER_CHOICES = {",": ",", ";": ";", "|": "|", "TAB": "\t", "SPACE": " "}
 
 
+def _openfile(filename, mode):
+    if filename.endswith(".gz"):
+        return gzip.open(filename, mode)
+    elif filename.endswith(".bz2"):
+        return bz2.open(filename, mode)
+    elif filename.endswith(".xz"):
+        return lzma.open(filename, mode)
+    elif "b" in mode:
+        return open(filename, mode)  # pylint: disable=unspecified-encoding
+    else:
+        return open(filename, mode, encoding="utf-8")
+
+
 class LoadCsvPlugin(LoadFileDeGlobMixin, LoadFileWithFilenameMixin, DuckdbLoadFilePlugin):
     """Load CSV files"""
 
@@ -70,29 +83,19 @@ class LoadCsvPlugin(LoadFileDeGlobMixin, LoadFileWithFilenameMixin, DuckdbLoadFi
     def load_file(
         self, cursor: duckdb.DuckDBPyConnection, filename: str, file_param: BaseParam, row_limit: Optional[int] = None
     ) -> duckdb.DuckDBPyRelation:
-        options = {
-            "sep": CSV_DELIMITER_CHOICES[self.delimiter.value],
-            "null_padding": True,
-            "strict_mode": False,
-        }
-
-        if len(self.columns):
-            options["all_varchar"] = True
-            options["skiprows"] = 1 if self.header else 0
-            options["header"] = False
-        else:
-            options["header"] = self.header.value
-
-        if filename.endswith(".xz"):
-            logger.debug("Reading file %s with LZMA", filename)
-            with lzma.open(filename) as fh:
-                rel = duckdb_source_to_table(cursor, cursor.read_csv(fh, **options))
-        if filename.endswith(".bz2"):
-            logger.debug("Reading file %s with BZ2", filename)
-            with bz2.open(filename) as fh:
-                rel = duckdb_source_to_table(cursor, cursor.read_csv(fh, **options))
-        else:
-            rel = duckdb_source_to_view(cursor, cursor.read_csv(filename, **options))
+        with _openfile(filename, "rb") as fh:
+            rel = duckdb_source_to_table(
+                cursor,
+                cursor.read_csv(
+                    fh,
+                    sep=CSV_DELIMITER_CHOICES[self.delimiter.value],
+                    null_padding=True,
+                    strict_mode=False,
+                    all_varchar=len(self.columns) > 0,
+                    skiprows=1 if self.header and len(self.columns) > 0 else 0,
+                    header=self.header.value and len(self.columns) == 0,
+                ),
+            )
 
         if row_limit is not None:
             rel = rel.limit(row_limit)
@@ -119,7 +122,6 @@ def load_file(
                 if cp is None or cp.type.is_not_none()
             )
 
-
             logger.debug("LoadCsvPlugin.load_file proj %s", proj)
             rel = rel.project(proj)
 
@@ -174,7 +176,6 @@ class SaveCsvPlugin(DuckdbSaveFilePlugin):
     def execute(
         self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation], row_limit: Optional[int] = None
     ) -> None:
-
         if source is None or row_limit is not None or not self.filename:
             return
 
@@ -185,7 +186,6 @@ def execute(
             logger.debug("SaveCsvPlugin.execute order_by %s", order_by)
             source = source.order(order_by)
 
-
         options = {
             "index": False,
             "sep": self.SEPARATORS[self.delimiter.value],
@@ -195,20 +195,12 @@ def execute(
         filename = self.filename.get_file_path()
         Path(filename).parent.mkdir(parents=True, exist_ok=True)
 
-        def _openfile(filename):
-            if filename.endswith(".gz"):
-                return gzip.open(filename, "wb")
-            elif filename.endswith(".bz2"):
-                return bz2.open(filename, "wb")
-            else:
-                return open(filename, "wb")
-
         # the type check supression is because the first parameter to
         # to_csv is called path_or_buf and takes either a filename path
         # or a file-like buffer, but is declared as str|None.
 
         logger.debug("writing file %s", filename)
-        with _openfile(filename) as fh:
+        with _openfile(filename, "wb") as fh:
             chunk = source.fetch_df_chunk()
             chunk.to_csv(fh, header=True, **options)  # type: ignore
             while len(chunk) > 0:

countess/utils/expression_to_sql.py

@@ -194,6 +194,7 @@ class AddExpr(SqlTemplatingList):
 class CompOp(SqlTemplatingSymbol):
     regex = re.compile(r"<=?|>=?|[!=]=")
 
+
 class CompExpr(pypeg2.List):
     grammar = AddExpr, pypeg2.maybe_some((CompOp, [NullLiteral, AddExpr]))
 
@@ -208,10 +209,7 @@ def _subcomp(left, comp, right):
             return left.sql() + comp_op + right.sql()
 
         if len(self) >= 3:
-            return " AND ".join(
-                _subcomp(self[n], self[n+1], self[n+2])
-                for n in range(0, len(self)-1, 2)
-            )
+            return " AND ".join(_subcomp(self[n], self[n + 1], self[n + 2]) for n in range(0, len(self) - 1, 2))
         else:
             return self[0].sql()
 
@@ -256,10 +254,7 @@ def names(self):
         return [s.name for s in self[0:-1]]
 
     def sql(self):
-        return ','.join(
-            self[-1].sql() + " AS " + s.sql()
-            for s in self[0:-1]
-        )
+        return ",".join(self[-1].sql() + " AS " + s.sql() for s in self[0:-1])
 
 
 class Filter(pypeg2.Concat):

tests/test_expression.py

@@ -26,27 +26,32 @@
 def test_expressions(expr, sql):
     assert Block.from_string(expr)[0].sql() == sql
 
+
 assign_tests = [
     ("a = 1", '1 AS "a"'),
     ("a = b = c = 2", '2 AS "a",2 AS "b",2 AS "c"'),
     ("d = None", 'NULL AS "d"'),
 ]
 
+
 @pytest.mark.parametrize("expr,sql", assign_tests)
 def test_assignments(expr, sql):
     assert Block.from_string(expr).sql_selects()[0] == sql
 
+
 filter_tests = [
     ("a < b", '"a"<"b"'),
     ("b <= c", '"b"<="c"'),
     ("a < b <= c", '"a"<"b" AND "b"<="c"'),
     ("__filter = a < b", '"a"<"b"'),
 ]
 
+
 @pytest.mark.parametrize("expr,sql", filter_tests)
 def test_filters(expr, sql):
     assert Block.from_string(expr).sql_wheres()[0] == sql
 
+
 def test_project_and_filter():
     rel = duckdb.sql("select * from (values (1,2), (3,4), (5,6)) t(a,b)")
     blk = Block.from_string("c = a+b\nb<5\nd = a + c")