Flowfile 0.4.1__py3-none-any.whl → 0.5.3__py3-none-any.whl

flowfile_worker/create/funcs.py

@@ -1,95 +1,129 @@
-import polars as pl
 import os
 
-from flowfile_worker.create.models import ReceivedCsvTable, ReceivedParquetTable, ReceivedExcelTable
+import polars as pl
+
+from flowfile_worker.create.models import (
+    InputCsvTable,
+    InputExcelTable,
+    InputJsonTable,
+    InputParquetTable,
+    ReceivedTable,
+)
+from flowfile_worker.create.read_excel_tables import df_from_calamine_xlsx, df_from_openpyxl
 from flowfile_worker.create.utils import create_fake_data
-from flowfile_worker.create.read_excel_tables import df_from_openpyxl, df_from_calamine_xlsx
 
 
-def create_from_path_json(received_table: ReceivedCsvTable):
+def create_from_path_json(received_table: ReceivedTable):
+    if not isinstance(received_table.table_settings, InputJsonTable):
+        raise ValueError("Received table settings are not of type InputJsonTable")
+    input_table_settings: InputJsonTable = received_table.table_settings
     f = received_table.abs_file_path
     gbs_to_load = os.path.getsize(f) / 1024 / 1000 / 1000
     low_mem = gbs_to_load > 10
-    if received_table.encoding.upper() == 'UTF8' or received_table.encoding.upper() == 'UTF-8':
+    if input_table_settings.encoding.upper() == "UTF8" or input_table_settings.encoding.upper() == "UTF-8":
         try:
-            df = pl.scan_csv(f,
-                               low_memory=low_mem,
-                               try_parse_dates=True,
-                               separator=received_table.delimiter,
-                               has_header=received_table.has_headers,
-                               skip_rows=received_table.starting_from_line,
-                               encoding='utf8',
-                               infer_schema_length=received_table.infer_schema_length)
+            df = pl.scan_csv(
+                f,
+                low_memory=low_mem,
+                try_parse_dates=True,
+                separator=input_table_settings.delimiter,
+                has_header=input_table_settings.has_headers,
+                skip_rows=input_table_settings.starting_from_line,
+                encoding="utf8",
+                infer_schema_length=input_table_settings.infer_schema_length,
+            )
             df.head(1).collect()
             return df
         except:
             try:
-                df = pl.scan_csv(f, low_memory=low_mem,
-                                   separator=received_table.delimiter,
-                                   has_header=received_table.has_headers,
-                                   skip_rows=received_table.starting_from_line,
-                                   encoding='utf8-lossy',
-                                   ignore_errors=True)
+                df = pl.scan_csv(
+                    f,
+                    low_memory=low_mem,
+                    separator=input_table_settings.delimiter,
+                    has_header=input_table_settings.has_headers,
+                    skip_rows=input_table_settings.starting_from_line,
+                    encoding="utf8-lossy",
+                    ignore_errors=True,
+                )
                 return df
             except:
-                df = pl.scan_csv(f, low_memory=low_mem,
-                                   separator=received_table.delimiter,
-                                   has_header=received_table.has_headers,
-                                   skip_rows=received_table.starting_from_line,
-                                   encoding='utf8',
-                                   ignore_errors=True)
+                df = pl.scan_csv(
+                    f,
+                    low_memory=low_mem,
+                    separator=input_table_settings.delimiter,
+                    has_header=input_table_settings.has_headers,
+                    skip_rows=input_table_settings.starting_from_line,
+                    encoding="utf8",
+                    ignore_errors=True,
+                )
                 return df
     else:
-        df = pl.read_csv(f, low_memory=low_mem,
-                           separator=received_table.delimiter,
-                           has_header=received_table.has_headers,
-                           skip_rows=received_table.starting_from_line,
-                           encoding=received_table.encoding,
-                           ignore_errors=True)
+        df = pl.read_csv(
+            f,
+            low_memory=low_mem,
+            separator=input_table_settings.delimiter,
+            has_header=input_table_settings.has_headers,
+            skip_rows=input_table_settings.starting_from_line,
+            encoding=input_table_settings.encoding,
+            ignore_errors=True,
+        )
         return df
 
 
-def create_from_path_csv(received_table: ReceivedCsvTable) -> pl.DataFrame:
+def create_from_path_csv(received_table: ReceivedTable) -> pl.DataFrame:
     f = received_table.abs_file_path
+    if not isinstance(received_table.table_settings, InputCsvTable):
+        raise ValueError("Received table settings are not of type InputCsvTable")
+    input_table_settings: InputCsvTable = received_table.table_settings
     gbs_to_load = os.path.getsize(f) / 1024 / 1000 / 1000
     low_mem = gbs_to_load > 10
-    if received_table.encoding.upper() == 'UTF8' or received_table.encoding.upper() == 'UTF-8':
+    if input_table_settings.encoding.upper() == "UTF8" or input_table_settings.encoding.upper() == "UTF-8":
         try:
-            df = pl.scan_csv(f,
-                               low_memory=low_mem,
-                               try_parse_dates=True,
-                               separator=received_table.delimiter,
-                               has_header=received_table.has_headers,
-                               skip_rows=received_table.starting_from_line,
-                               encoding='utf8',
-                               infer_schema_length=received_table.infer_schema_length)
+            df = pl.scan_csv(
+                f,
+                low_memory=low_mem,
+                try_parse_dates=True,
+                separator=input_table_settings.delimiter,
+                has_header=input_table_settings.has_headers,
+                skip_rows=input_table_settings.starting_from_line,
+                encoding="utf8",
+                infer_schema_length=input_table_settings.infer_schema_length,
+            )
             df.head(1).collect()
             return df
         except:
             try:
-                df = pl.scan_csv(f, low_memory=low_mem,
-                                   separator=received_table.delimiter,
-                                   has_header=received_table.has_headers,
-                                   skip_rows=received_table.starting_from_line,
-                                   encoding='utf8-lossy',
-                                   ignore_errors=True)
+                df = pl.scan_csv(
+                    f,
+                    low_memory=low_mem,
+                    separator=input_table_settings.delimiter,
+                    has_header=input_table_settings.has_headers,
+                    skip_rows=input_table_settings.starting_from_line,
+                    encoding="utf8-lossy",
+                    ignore_errors=True,
+                )
                 return df
             except:
-                df = pl.scan_csv(f, low_memory=low_mem,
-                                   separator=received_table.delimiter,
-                                   has_header=received_table.has_headers,
-                                   skip_rows=received_table.starting_from_line,
-                                   encoding='utf8',
-                                   ignore_errors=True)
+                df = pl.scan_csv(
+                    f,
+                    low_memory=low_mem,
+                    separator=input_table_settings.delimiter,
+                    has_header=input_table_settings.has_headers,
+                    skip_rows=input_table_settings.starting_from_line,
+                    encoding="utf8",
+                    ignore_errors=True,
+                )
                 return df
     else:
-        df = pl.read_csv(f,
-                           low_memory=low_mem,
-                           separator=received_table.delimiter,
-                           has_header=received_table.has_headers,
-                           skip_rows=received_table.starting_from_line,
-                           encoding=received_table.encoding,
-                           ignore_errors=True)
+        df = pl.read_csv(
+            f,
+            low_memory=low_mem,
+            separator=input_table_settings.delimiter,
+            has_header=input_table_settings.has_headers,
+            skip_rows=input_table_settings.starting_from_line,
+            encoding=input_table_settings.encoding,
+            ignore_errors=True,
+        )
         return df
 
 
@@ -97,50 +131,65 @@ def create_random(number_of_records: int = 1000) -> pl.LazyFrame:
     return create_fake_data(number_of_records).lazy()
 
 
-def create_from_path_parquet(received_table: ReceivedParquetTable):
+def create_from_path_parquet(received_table: ReceivedTable):
+    if not isinstance(received_table.table_settings, InputParquetTable):
+        raise ValueError("Received table settings are not of type InputParquetTable")
     low_mem = (os.path.getsize(received_table.abs_file_path) / 1024 / 1000 / 1000) > 2
     return pl.scan_parquet(source=received_table.abs_file_path, low_memory=low_mem)
 
 
-def create_from_path_excel(received_table: ReceivedExcelTable):
-    if received_table.type_inference:
-        engine = 'openpyxl'
-    elif received_table.start_row > 0 and received_table.start_column == 0:
-        engine = 'calamine' if received_table.has_headers else 'xlsx2csv'
-    elif received_table.start_column > 0 or received_table.start_row > 0:
-        engine = 'openpyxl'
+def create_from_path_excel(received_table: ReceivedTable):
+    if not isinstance(received_table.table_settings, InputExcelTable):
+        raise ValueError("Received table settings are not of type InputExcelTable")
+    input_table_settings: InputExcelTable = received_table.table_settings
+
+    if input_table_settings.type_inference:
+        engine = "openpyxl"
+    elif input_table_settings.start_row > 0 and input_table_settings.start_column == 0:
+        engine = "calamine" if input_table_settings.has_headers else "xlsx2csv"
+    elif input_table_settings.start_column > 0 or input_table_settings.start_row > 0:
+        engine = "openpyxl"
     else:
-        engine = 'calamine'
+        engine = "calamine"
 
-    sheet_name = received_table.sheet_name
+    sheet_name = input_table_settings.sheet_name
 
-    if engine == 'calamine':
-        df = df_from_calamine_xlsx(file_path=received_table.abs_file_path, sheet_name=sheet_name,
-                                   start_row=received_table.start_row, end_row=received_table.end_row)
-        if received_table.end_column > 0:
-            end_col_index = received_table.end_column
-            cols_to_select = [df.columns[i] for i in range(received_table.start_column, end_col_index)]
+    if engine == "calamine":
+        df = df_from_calamine_xlsx(
+            file_path=received_table.abs_file_path,
+            sheet_name=sheet_name,
+            start_row=input_table_settings.start_row,
+            end_row=input_table_settings.end_row,
+        )
+        if input_table_settings.end_column > 0:
+            end_col_index = input_table_settings.end_column
+            cols_to_select = [df.columns[i] for i in range(input_table_settings.start_column, end_col_index)]
             df = df.select(cols_to_select)
 
-    elif engine == 'xlsx2csv':
-        csv_options = {'has_header': received_table.has_headers, 'skip_rows': received_table.start_row}
-        df = pl.read_excel(source=received_table.abs_file_path,
-                           read_options=csv_options,
-                           engine='xlsx2csv',
-                           sheet_name=received_table.sheet_name)
-        end_col_index = received_table.end_column if received_table.end_column > 0 else len(df.columns)
-        cols_to_select = [df.columns[i] for i in range(received_table.start_column, end_col_index)]
+    elif engine == "xlsx2csv":
+        csv_options = {"has_header": input_table_settings.has_headers, "skip_rows": input_table_settings.start_row}
+        df = pl.read_excel(
+            source=received_table.abs_file_path,
+            read_options=csv_options,
+            engine="xlsx2csv",
+            sheet_name=input_table_settings.sheet_name,
+        )
+        end_col_index = input_table_settings.end_column if input_table_settings.end_column > 0 else len(df.columns)
+        cols_to_select = [df.columns[i] for i in range(input_table_settings.start_column, end_col_index)]
         df = df.select(cols_to_select)
-        if 0 < received_table.end_row < len(df):
-            df = df.head(received_table.end_row)
+        if 0 < input_table_settings.end_row < len(df):
+            df = df.head(input_table_settings.end_row)
 
     else:
-        max_col = received_table.end_column if received_table.end_column > 0 else None
-        max_row = received_table.end_row + 1 if received_table.end_row > 0 else None
-        df = df_from_openpyxl(file_path=received_table.abs_file_path,
-                              sheet_name=received_table.sheet_name,
-                              min_row=received_table.start_row + 1,
-                              min_col=received_table.start_column + 1,
-                              max_row=max_row,
-                              max_col=max_col, has_headers=received_table.has_headers)
+        max_col = input_table_settings.end_column if input_table_settings.end_column > 0 else None
+        max_row = input_table_settings.end_row + 1 if input_table_settings.end_row > 0 else None
+        df = df_from_openpyxl(
+            file_path=received_table.abs_file_path,
+            sheet_name=input_table_settings.sheet_name,
+            min_row=input_table_settings.start_row + 1,
+            min_col=input_table_settings.start_column + 1,
+            max_row=max_row,
+            max_col=max_col,
+            has_headers=input_table_settings.has_headers,
+        )
     return df

flowfile_worker/create/models.py

@@ -1,7 +1,8 @@
-from pydantic import BaseModel, Field, model_validator
-from typing import List, Optional
 import os
 from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field, field_validator, model_validator
 
 
 class MinimalFieldInfo(BaseModel):
@@ -9,78 +10,148 @@ class MinimalFieldInfo(BaseModel):
     data_type: str
 
 
-class ReceivedTableBase(BaseModel):
-    id: Optional[int] = None
-    name: str
-    path: str
-    directory: Optional[str] = None
-    analysis_file_available: Optional[bool] = False
-    status: Optional[str] = None
-    file_type: Optional[str] = None
-    fields: List[MinimalFieldInfo] = Field(default_factory=list)
-    abs_file_path: Optional[str] = None
+class InputTableBase(BaseModel):
+    """Base settings for input file operations."""
 
-    @classmethod
-    def create_from_path(cls, path: str):
-        filename = os.path.basename(path)
-        return cls(name=filename, path=path)
+    file_type: str  # Will be overridden with Literal in subclasses
 
-    @property
-    def file_path(self) -> str:
-        if self.name not in self.path:
-            return os.path.join(self.path, self.name)
-        return self.path
 
-    @model_validator(mode="after")
-    def set_abs_file_path(cls, values):
-        abs_file_path = getattr(values, "abs_file_path", None)
-        if abs_file_path is None:
-            path = getattr(values, "path", None)
-            if not path:
-                raise ValueError("Field 'path' is required to compute abs_file_path")
-            setattr(values, "abs_file_path", str(Path(path).absolute()))
-        return values
-
-
-class ReceivedCsvTable(ReceivedTableBase):
-    file_type: Optional[str] = 'csv'
-    reference: Optional[str] = ''
-    starting_from_line: Optional[int] = 0
-    delimiter: Optional[str] = ','
-    has_headers: Optional[bool] = True
-    encoding: Optional[str] = 'utf-8'
-    parquet_ref: Optional[str] = None
-    row_delimiter: Optional[str] = '\n'
-    quote_char: Optional[str] = '"'
-    infer_schema_length: Optional[int] = 10_000
-    truncate_ragged_lines: Optional[bool] = False
-    ignore_errors: Optional[bool] = False
-
-
-class ReceivedJsonTable(ReceivedCsvTable):
-    pass
-
-
-class ReceivedParquetTable(ReceivedTableBase):
-    file_type: Optional[str] = 'parquet'
-
-
-class ReceivedExcelTable(ReceivedTableBase):
-    sheet_name: Optional[str] = None
-    start_row: Optional[int] = 0  # optional
-    start_column: Optional[int] = 0  # optional
-    end_row: Optional[int] = 0  # optional
-    end_column: Optional[int] = 0  # optional
-    has_headers: Optional[bool] = True  # optional
-    type_inference: Optional[bool] = False  # optional
+class InputCsvTable(InputTableBase):
+    """Defines settings for reading a CSV file."""
+
+    file_type: Literal["csv"] = "csv"
+    reference: str = ""
+    starting_from_line: int = 0
+    delimiter: str = ","
+    has_headers: bool = True
+    encoding: str = "utf-8"
+    parquet_ref: str | None = None
+    row_delimiter: str = "\n"
+    quote_char: str = '"'
+    infer_schema_length: int = 10_000
+    truncate_ragged_lines: bool = False
+    ignore_errors: bool = False
+
+
+class InputJsonTable(InputCsvTable):
+    """Defines settings for reading a JSON file."""
+
+    file_type: Literal["json"] = "json"
+
+
+class InputParquetTable(InputTableBase):
+    """Defines settings for reading a Parquet file."""
+
+    file_type: Literal["parquet"] = "parquet"
+
+
+class InputExcelTable(InputTableBase):
+    """Defines settings for reading an Excel file."""
+
+    file_type: Literal["excel"] = "excel"
+    sheet_name: str | None = None
+    start_row: int = 0
+    start_column: int = 0
+    end_row: int = 0
+    end_column: int = 0
+    has_headers: bool = True
+    type_inference: bool = False
 
+    @model_validator(mode="after")
     def validate_range_values(self):
-        # Validate that start and end rows/columns are non-negative integers
+        """Validates that the Excel cell range is logical."""
         for attribute in [self.start_row, self.start_column, self.end_row, self.end_column]:
             if not isinstance(attribute, int) or attribute < 0:
                 raise ValueError("Row and column indices must be non-negative integers")
+        if (self.end_row > 0 and self.start_row > self.end_row) or (
+            self.end_column > 0 and self.start_column > self.end_column
+        ):
+            raise ValueError("Start row/column must not be greater than end row/column")
+        return self
+
+
+# Create the discriminated union (similar to OutputTableSettings)
+InputTableSettings = Annotated[
+    InputCsvTable | InputJsonTable | InputParquetTable | InputExcelTable, Field(discriminator="file_type")
+]
+
+
+# Now create the main ReceivedTable model
+class ReceivedTable(BaseModel):
+    """Model for defining a table received from an external source."""
+
+    # Metadata fields
+    id: int | None = None
+    name: str | None = None
+    path: str  # This can be an absolute or relative path
+    directory: str | None = None
+    analysis_file_available: bool = False
+    status: str | None = None
+    fields: list[MinimalFieldInfo] = Field(default_factory=list)
+    abs_file_path: str | None = None
+
+    file_type: Literal["csv", "json", "parquet", "excel"]
+
+    table_settings: InputTableSettings
 
-        # Validate that start is before end if end is specified (non-zero)
-        if (0 < self.end_row < self.start_row) or \
-                (0 < self.end_column < self.start_column):
-            raise ValueError("Start row/column must not be greater than end row/column if specified")
+    @classmethod
+    def create_from_path(cls, path: str, file_type: Literal["csv", "json", "parquet", "excel"] = "csv"):
+        """Creates an instance from a file path string."""
+        filename = Path(path).name
+
+        # Create appropriate table_settings based on file_type
+        settings_map = {
+            "csv": InputCsvTable(),
+            "json": InputJsonTable(),
+            "parquet": InputParquetTable(),
+            "excel": InputExcelTable(),
+        }
+
+        return cls(
+            name=filename, path=path, file_type=file_type, table_settings=settings_map.get(file_type, InputCsvTable())
+        )
+
+    @property
+    def file_path(self) -> str:
+        """Constructs the full file path from the directory and name."""
+        if self.name and self.name not in self.path:
+            return os.path.join(self.path, self.name)
+        else:
+            return self.path
+
+    def set_absolute_filepath(self):
+        """Resolves the path to an absolute file path."""
+        base_path = Path(self.path).expanduser()
+        if not base_path.is_absolute():
+            base_path = Path.cwd() / base_path
+        if self.name and self.name not in base_path.name:
+            base_path = base_path / self.name
+        self.abs_file_path = str(base_path.resolve())
+
+    @field_validator("table_settings", mode="before")
+    @classmethod
+    def validate_table_settings(cls, v, info):
+        """Ensures table_settings matches the file_type."""
+        if v is None:
+            file_type = info.data.get("file_type", "csv")
+            # Create default based on file_type
+            settings_map = {
+                "csv": InputCsvTable(),
+                "json": InputJsonTable(),
+                "parquet": InputParquetTable(),
+                "excel": InputExcelTable(),
+            }
+            return settings_map.get(file_type, InputCsvTable())
+
+        # If it's a dict, add file_type if missing
+        if isinstance(v, dict) and "file_type" not in v:
+            v["file_type"] = info.data.get("file_type", "csv")
+
+        return v
+
+    @model_validator(mode="after")
+    def populate_abs_file_path(self):
+        """Ensures the absolute file path is populated after validation."""
+        if not self.abs_file_path:
+            self.set_absolute_filepath()
+        return self

flowfile_worker/create/pl_types.py

@@ -1,21 +1,20 @@
 import polars as pl
 
-
 dtype_to_pl = {
-    'int': pl.Int64,
-    'integer': pl.Int64,
-    'char': pl.String,
-    'fixed decimal': pl.Float32,
-    'double': pl.Float64,
-    'float': pl.Float64,
-    'bool': pl.Boolean,
-    'byte': pl.UInt8,
-    'bit': pl.Binary,
-    'date': pl.Date,
-    'datetime': pl.Datetime,
-    'string': pl.String,
-    'str': pl.String,
-    'time': pl.Time,
+    "int": pl.Int64,
+    "integer": pl.Int64,
+    "char": pl.String,
+    "fixed decimal": pl.Float32,
+    "double": pl.Float64,
+    "float": pl.Float64,
+    "bool": pl.Boolean,
+    "byte": pl.UInt8,
+    "bit": pl.Binary,
+    "date": pl.Date,
+    "datetime": pl.Datetime,
+    "string": pl.String,
+    "str": pl.String,
+    "time": pl.Time,
 }
 
 dtype_to_pl_str = {k: v.__name__ for k, v in dtype_to_pl.items()}