Alex Galarce
0ad2798d6b
fix quoting/encoding issue in crosswalk writing to CSV (also formatting) * fix encoding issue * black formatting * isort Approved-by: Katon Minhas
146 lines
4.4 KiB
Python
146 lines
4.4 KiB
Python
import csv
|
|
import json
|
|
|
|
import pandas as pd
|
|
|
|
|
|
class CrosswalkBuilder:
|
|
def __init__(self):
|
|
self.mapping: dict[str, str] = {}
|
|
self.metadata: dict[str, str | None] = {}
|
|
|
|
def from_excel(
|
|
self, path: str, from_col: str, to_col: str, sheet_name: str | None = None
|
|
) -> "CrosswalkBuilder":
|
|
"""Load mapping from an excel file
|
|
|
|
Args:
|
|
path (str): location for the excel file
|
|
from_col (str): Source column (for keys)
|
|
to_col (str): Destination column (for values)
|
|
sheet_name (str|None, optional): Sheet name in the excel file. Defaults to None.
|
|
|
|
Returns:
|
|
CrosswalkBuilder: the builder instance for method chaining
|
|
"""
|
|
if sheet_name is None:
|
|
# If sheet_name is not provided, read the first sheet
|
|
df = pd.read_excel(path, sheet_name=0)
|
|
else:
|
|
df = pd.read_excel(path, sheet_name=sheet_name)
|
|
|
|
if from_col not in df.columns:
|
|
raise ValueError(f"Column '{from_col}' not found in the excel file")
|
|
elif to_col not in df.columns:
|
|
raise ValueError(f"Column '{to_col}' not found in the excel file")
|
|
|
|
self.metadata = {
|
|
"from_col": from_col,
|
|
"to_col": to_col,
|
|
}
|
|
|
|
self.mapping = dict(zip(df[from_col], df[to_col]))
|
|
|
|
return self
|
|
|
|
def from_dict(
|
|
self,
|
|
mapping: dict[str, str],
|
|
from_col: str | None = None,
|
|
to_col: str | None = None,
|
|
) -> "CrosswalkBuilder":
|
|
"""Load mapping from a dictionary
|
|
|
|
Args:
|
|
mapping (dict[str, str]): Mapping dictionary
|
|
|
|
Returns:
|
|
CrosswalkBuilder: the builder instance for method chaining
|
|
"""
|
|
self.metadata = {
|
|
"from_col": from_col,
|
|
"to_col": to_col,
|
|
}
|
|
self.mapping = mapping
|
|
return self
|
|
|
|
def from_json(self, path: str) -> "CrosswalkBuilder":
|
|
"""Load mapping from a JSON file
|
|
|
|
Args:
|
|
path (str): Path to the JSON file
|
|
|
|
Returns:
|
|
CrosswalkBuilder: the builder instance for method chaining
|
|
"""
|
|
with open(path, "r") as f:
|
|
data = json.load(f)
|
|
|
|
self.metadata = data.get("metadata", {})
|
|
self.mapping = data.get("mapping", {})
|
|
|
|
return self
|
|
|
|
def save(self, output_path: str) -> None:
|
|
"""Save the crosswalk to JSON.
|
|
|
|
Args:
|
|
output_path (str): Path to save the crosswalk
|
|
"""
|
|
output_dict = {"metadata": self.metadata, "mapping": self.mapping}
|
|
with open(output_path, "w") as f:
|
|
json.dump(output_dict, f, indent=4)
|
|
|
|
def to_csv(self, output_path: str) -> None:
|
|
"""Save the crosswalk to CSV.
|
|
|
|
Args:
|
|
output_path (str): Path to save the crosswalk
|
|
"""
|
|
from_col = (
|
|
self.metadata.get("from_col") if self.metadata.get("from_col") else "source"
|
|
)
|
|
to_col = (
|
|
self.metadata.get("to_col") if self.metadata.get("to_col") else "target"
|
|
)
|
|
|
|
df = pd.DataFrame(self.mapping.items(), columns=[from_col, to_col])
|
|
df.to_csv(output_path, index=False, quoting=csv.QUOTE_ALL, encoding="utf-8-sig")
|
|
|
|
def create_reverse_mapping(self) -> dict[str, str]:
|
|
"""Create a reverse mapping where values that map to the same key are concatenated
|
|
|
|
Example:
|
|
{"A": "1", "B": "1", "C": "2"} becomes {"1": "A, B", "2": "C"}
|
|
|
|
Returns:
|
|
dict[str, str]: Reverse mapping with concatenated values
|
|
"""
|
|
reverse_dict: dict[str, list[str]] = {}
|
|
|
|
# First collect all sources for each target
|
|
for source, target in self.mapping.items():
|
|
if target in reverse_dict:
|
|
reverse_dict[target].append(source)
|
|
else:
|
|
reverse_dict[target] = [source]
|
|
|
|
# Then join the lists with commas
|
|
return {target: ", ".join(sources) for target, sources in reverse_dict.items()}
|
|
|
|
|
|
def apply_crosswalk(
|
|
val: str, mapping: dict[str, str], default: str | None = None
|
|
) -> str:
|
|
"""Apply a crosswalk to a value
|
|
|
|
Args:
|
|
val (str): Input value
|
|
mapping (dict[str, str]): Mapping dictionary
|
|
default (str | None, optional): Default value; if None and val is not found in the mapping, val is returned. Defaults to None.
|
|
|
|
Returns:
|
|
str: Crosswalked value
|
|
"""
|
|
return mapping.get(val, default if default is not None else val)
|