stringify list columns and add zstandard #97

Author: HongPong

.config/mise.toml

@@ -7,3 +7,4 @@ actionlint = "latest"
 shellcheck = "latest"
 markdownlint-cli2 = "latest"
 jq = "latest"
+uv = "latest"

file_utils.py

@@ -1,6 +1,7 @@
 import copy
 import json
 import os
+import polars as pl
 from schemas import enrichment_print_schema
 from utils import (
     convert_to_dataframe,
@@ -9,6 +10,15 @@
 )
 import xlsxwriter  # type: ignore [import-untyped]
 
+# Deals with list columns data that CSV cannot deal with.
+def _stringify_list_columns(df: pl.DataFrame) -> pl.DataFrame:
+    """Convert any List-type columns to JSON strings so CSV/Excel can handle them."""
+    list_cols = [col for col, dtype in zip(df.columns, df.dtypes) if dtype.base_type() == pl.List]
+    if list_cols:
+        df = df.with_columns(
+            [pl.col(c).map_elements(lambda val: json.dumps(val, default=str), return_dtype=pl.String).alias(c) for c in list_cols]
+        )
+    return df
 
 def export_to_file(
     facilities_data: dict,
@@ -24,10 +34,12 @@ def export_to_file(
         match file_type:
             case "xlsx":
                 with xlsxwriter.Workbook(full_name, {"remove_timezone": True}) as wb:
-                    _ = writer.write_excel(workbook=wb, include_header=True, autofit=True)
+                    _ = _stringify_list_columns(writer).write_excel(workbook=wb, include_header=True, autofit=True)
+                    # _ = writer.write_excel(workbook=wb, include_header=True, autofit=True)
             case "csv":
                 with open(full_name, "w", newline="", encoding="utf-8") as f_out:
-                    writer.write_csv(file=f_out, include_header=True)
+                    # writer.write_csv(file=f_out, include_header=True)
+                    _stringify_list_columns(writer).write_csv(file=f_out, include_header=True)
             case "parquet":
                 writer.write_parquet(full_name, use_pyarrow=True)
             case _:

ice_scrapers/inspections.py

@@ -1,5 +1,5 @@
 from bs4 import BeautifulSoup
-from compression import zstd
+import zstandard as zstd
 import os
 import pdfplumber
 from pprint import pformat

pyproject.toml

@@ -14,6 +14,7 @@ dependencies = [
     "requests>=2.32.5",
     "thefuzz>=0.22.1",
     "xlsxwriter>=3.2.5",
+    "zstandard>=0.25.0",
 ]
 
 [dependency-groups]