45 lines
1.7 KiB
Python
45 lines
1.7 KiB
Python
from __future__ import annotations
|
|
from typing import Dict
|
|
import polars as pl
|
|
from tools.base import BaseTool
|
|
|
|
|
|
class DateTimeTool(BaseTool):
|
|
def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
|
|
df = inputs.get("Input", pl.DataFrame())
|
|
if self.config is None or df.is_empty():
|
|
return {"Output": df}
|
|
|
|
fmt = self._cfg("Format", "%Y-%m-%d") or "%Y-%m-%d"
|
|
input_field = self._cfg("InputField", "") or ""
|
|
output_field = self._cfg("OutputField", input_field) or input_field
|
|
input_type = self._cfg("InputType", "String") or "String"
|
|
output_type = self._cfg("OutputType", "Date") or "Date"
|
|
|
|
if not input_field or input_field not in df.columns:
|
|
return {"Output": df}
|
|
|
|
# Convert Alteryx format to Python/Polars strptime format (they match mostly)
|
|
col = pl.col(input_field)
|
|
|
|
if input_type == "String":
|
|
if output_type == "Date":
|
|
result = col.str.to_date(fmt, strict=False)
|
|
elif output_type == "DateTime":
|
|
result = col.str.to_datetime(fmt, strict=False)
|
|
elif output_type == "Time":
|
|
result = col.str.to_time(fmt, strict=False)
|
|
else:
|
|
result = col.str.to_datetime(fmt, strict=False)
|
|
else:
|
|
# Date/DateTime → String
|
|
if input_type == "Date":
|
|
result = col.dt.strftime(fmt)
|
|
elif input_type == "DateTime":
|
|
result = col.dt.strftime(fmt)
|
|
else:
|
|
result = col.cast(pl.String)
|
|
|
|
df = df.with_columns(result.alias(output_field))
|
|
return {"Output": df}
|