dr-frames 0.1.0__py3-none-any.whl

dr_frames/filtering.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any, cast
+
+import pandas as pd
+
+__all__ = [
+    "select_subset",
+    "apply_filters_to_df",
+    "filter_to_value",
+    "filter_to_values",
+    "filter_to_range",
+    "filter_to_best_metric",
+    "make_filter_fxn",
+]
+
+
+def select_subset(
+    df: pd.DataFrame,
+    filters: Mapping[str, Any] | list[tuple[str, Any]] | None = None,
+) -> pd.DataFrame:
+    if filters is None:
+        filters = []
+    items = filters.items() if isinstance(filters, Mapping) else filters
+    mask = pd.Series(True, index=df.index)
+    for column, value in items:
+        assert column in df.columns, f"Column '{column}' not present in DataFrame."
+        if value is None or (isinstance(value, float) and pd.isna(value)):
+            mask &= df[column].isna()
+        else:
+            mask &= df[column] == value
+    return df.loc[mask].copy()
+
+
+def apply_filters_to_df(
+    df: pd.DataFrame, filters: dict[str, Sequence[Any]]
+) -> pd.DataFrame:
+    df = df.copy()
+    avail_cols = set(df.columns.tolist())
+    for k, v in filters.items():
+        if k not in avail_cols:
+            continue
+        df = cast(pd.DataFrame, df[df[k].isin(v)])
+    return df.reset_index(drop=True)
+
+
+def filter_to_value(
+    df: pd.DataFrame, column: str, value: float | str | None
+) -> pd.DataFrame:
+    """Filter to rows matching a specific value. Use None to match NaN values."""
+    if value is None:
+        return df[df[column].isna()].copy()
+    return df[df[column] == value].copy()
+
+
+def filter_to_values(
+    df: pd.DataFrame, column: str, values: list[float | str | None]
+) -> pd.DataFrame:
+    """Filter to rows matching any value in list. Use None in list to include NaN."""
+    if None in values:
+        other_values = [v for v in values if v is not None]
+        mask = df[column].isna() | df[column].isin(other_values)
+    else:
+        mask = df[column].isin(values)
+    return df[mask].copy()
+
+
+def filter_to_range(
+    df: pd.DataFrame, column: str, min_val: float, max_val: float
+) -> pd.DataFrame:
+    return df[(df[column] >= min_val) & (df[column] <= max_val)].copy()
+
+
+def filter_to_best_metric(
+    df: pd.DataFrame,
+    group_cols: list[str],
+    metric_col: str,
+    lower_is_better: bool = True,
+) -> pd.DataFrame:
+    if lower_is_better:
+        idx = df.groupby(group_cols)[metric_col].idxmin()
+    else:
+        idx = df.groupby(group_cols)[metric_col].idxmax()
+    return df.loc[idx].copy()
+
+
+def make_filter_fxn(
+    filters: list[tuple[Callable, ...]],
+) -> Callable[[pd.DataFrame], pd.DataFrame]:
+    def apply(df: pd.DataFrame) -> pd.DataFrame:
+        for fn, *args in filters:
+            df = fn(df, *args)
+        return df
+
+    return apply

dr_frames/formatting.py

@@ -0,0 +1,257 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    from rich.table import Table
+
+JustifyMethod = Literal["default", "left", "center", "right", "full"]
+
+__all__ = [
+    "format_table",
+    "format_coverage_table",
+    "FORMATTER_TYPES",
+    "OUTPUT_FORMATS",
+]
+
+OUTPUT_FORMATS = {
+    "console": "grid",
+    "markdown": "pipe",
+    "latex": "latex",
+    "plain": "plain",
+    "csv": "simple",
+}
+
+FORMATTER_TYPES: dict[str, Callable] = {
+    "scientific": lambda x, precision=2: f"{x:.{precision}e}"
+    if x is not None
+    else "None",
+    "decimal": lambda x, precision=3: f"{x:.{precision}f}" if x is not None else "None",
+    "integer": lambda x: f"{x:,.0f}" if x is not None else "None",
+    "comma": lambda x: f"{x:,}" if x is not None else "None",
+    "truncate": lambda x, max_length=50: (
+        str(x)[:max_length] + "..."
+        if x is not None and len(str(x)) > max_length
+        else (str(x) if x is not None else "None")
+    ),
+    "string": lambda x: str(x) if x is not None else "None",
+}
+
+COVERAGE_TABLE_CONFIG = {
+    "index": {"header": "#", "formatter": "integer"},
+    "column": {"header": "Column", "formatter": "truncate", "max_length": 35},
+    "coverage": {"header": "Coverage %", "formatter": "decimal", "precision": 1},
+}
+
+
+def format_table(
+    data: list[dict] | pd.DataFrame | list[list],
+    headers: list[str] | None = None,
+    output_format: str = "console",
+    column_config: dict[str, dict] | None = None,
+    title: str | None = None,
+    table_style: str = "lines",
+    disable_numparse: bool = True,
+) -> str | Table:
+    processed_data = _preprocess_data(data)
+    column_names = _get_column_names(data)
+    config = column_config or {}
+    formatted_data = _apply_column_formatting(processed_data, config, column_names)
+    final_headers = _resolve_headers(headers, column_names, config)
+
+    if output_format == "console":
+        return _create_rich_table(
+            formatted_data, final_headers, config, column_names, title, table_style
+        )
+    else:
+        from tabulate import tabulate
+
+        tablefmt = OUTPUT_FORMATS.get(output_format, "grid")
+        return tabulate(
+            formatted_data,
+            headers=final_headers,
+            tablefmt=tablefmt,
+            disable_numparse=disable_numparse,
+        )
+
+
+def format_coverage_table(
+    df: pd.DataFrame,
+    title: str = "Column Coverage",
+    output_format: str = "console",
+    table_style: str = "lines",
+    disable_numparse: bool = True,
+) -> str:
+    coverage_data = []
+    row_count = len(df) or df.shape[0]
+    for i, col in enumerate(df.columns):
+        if row_count == 0:
+            coverage = 0
+        else:
+            coverage = df[col].notna().sum() / row_count * 100
+        coverage_data.append({"index": i + 1, "column": col, "coverage": coverage})
+    result = f"{title} ({len(df.columns)} columns):\n"
+    table_result = format_table(
+        data=coverage_data,
+        output_format=output_format,
+        column_config=COVERAGE_TABLE_CONFIG,
+        table_style=table_style,
+        disable_numparse=disable_numparse,
+    )
+    result += str(table_result) if not isinstance(table_result, str) else table_result
+    return result
+
+
+def _preprocess_data(data: list[dict] | pd.DataFrame | list[list]) -> list[list]:
+    if isinstance(data, pd.DataFrame):
+        return data.to_numpy().tolist()
+    elif isinstance(data, list) and len(data) > 0:
+        if isinstance(data[0], dict):
+            # Build stable union of all keys: preserve first-row order, then append new keys
+            keys = list(data[0].keys())
+            keys_set = set(keys)
+            for row in data[1:]:
+                for key in row.keys():  # type: ignore[union-attr]
+                    if key not in keys_set:
+                        keys.append(key)
+                        keys_set.add(key)
+            return [[row.get(key) for key in keys] for row in data]  # type: ignore[union-attr]
+        else:
+            return list(data)  # type: ignore[arg-type]
+    return []
+
+
+def _get_column_names(data: list[dict] | pd.DataFrame | list[list]) -> list[str]:
+    if isinstance(data, pd.DataFrame):
+        return list(data.columns)
+    elif isinstance(data, list) and len(data) > 0 and isinstance(data[0], dict):
+        # Build stable union of all keys: preserve first-row order, then append new keys
+        keys = list(data[0].keys())
+        keys_set = set(keys)
+        for row in data[1:]:
+            for key in row.keys():  # type: ignore[union-attr]
+                if key not in keys_set:
+                    keys.append(key)
+                    keys_set.add(key)
+        return keys
+    else:
+        return []
+
+
+def _apply_column_formatting(
+    processed_data: list[list],
+    config: dict[str, dict],
+    column_names: list[str] | None = None,
+) -> list[list]:
+    if not processed_data:
+        return processed_data
+    formatted_data = []
+    for row in processed_data:
+        formatted_row = []
+        for col_idx, value in enumerate(row):
+            col_name = (
+                column_names[col_idx]
+                if column_names and col_idx < len(column_names)
+                else None
+            )
+            formatted_row.append(_format_value(value, col_name, config))
+        formatted_data.append(formatted_row)
+    return formatted_data
+
+
+def _format_value(
+    value: Any,
+    col_name: str | None,
+    config: dict[str, dict],
+) -> str:
+    if col_name and col_name in config:
+        col_config = config[col_name]
+        formatter_name = col_config.get("formatter", "string")
+        formatter = FORMATTER_TYPES.get(formatter_name, FORMATTER_TYPES["string"])
+        formatter_kwargs = {
+            k: v for k, v in col_config.items() if k not in ["header", "formatter"]
+        }
+        try:
+            return formatter(value, **formatter_kwargs)
+        except (TypeError, ValueError):
+            return str(value) if value is not None else "None"
+    return str(value) if value is not None else "None"
+
+
+def _resolve_headers(
+    headers: list[str] | None, column_names: list[str], config: dict[str, dict]
+) -> list[str]:
+    if headers is not None:
+        return headers
+    if config and column_names:
+        result_headers = []
+        for col_name in column_names:
+            if col_name in config:
+                result_headers.append(config[col_name].get("header", col_name))
+            else:
+                result_headers.append(col_name)
+        return result_headers
+    if column_names:
+        return column_names
+    return []
+
+
+def _create_rich_table(
+    formatted_data: list[list],
+    headers: list[str],
+    config: dict[str, dict],
+    column_names: list[str],
+    title: str | None = None,
+    table_style: str = "lines",
+) -> Table:
+    from rich.table import Table
+
+    if table_style == "zebra":
+        table = Table(
+            title=title,
+            show_header=True,
+            header_style="bold magenta",
+            row_styles=["", "dim"],
+        )
+    else:
+        table = Table(
+            title=title, show_header=True, header_style="bold magenta", show_lines=True
+        )
+
+    for i, header in enumerate(headers):
+        col_name = _get_column_name_for_index(column_names, i)
+        col_config = config.get(col_name, {})
+        justify = _get_rich_justify(col_config)
+        style = _get_rich_style(col_config)
+        table.add_column(header, justify=justify, style=style)
+
+    for row in formatted_data:
+        table.add_row(*[str(cell) for cell in row])
+
+    return table
+
+
+def _get_column_name_for_index(column_names: list[str], index: int) -> str:
+    return column_names[index] if index < len(column_names) else f"col_{index}"
+
+
+def _get_rich_justify(col_config: dict) -> JustifyMethod:
+    formatter = col_config.get("formatter", "string")
+    if formatter in ["scientific", "decimal", "integer", "comma"]:
+        return "right"
+    return "left"
+
+
+def _get_rich_style(col_config: dict) -> str | None:
+    formatter = col_config.get("formatter", "string")
+    style_map = {
+        "scientific": "yellow",
+        "decimal": "green",
+        "integer": "cyan",
+        "comma": "cyan",
+        "truncate": "dim",
+    }
+    return style_map.get(formatter)

dr_frames/parsing.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import ast
+
+import pandas as pd
+
+__all__ = [
+    "parse_first_element",
+    "sum_list_elements",
+    "is_homogeneous",
+    "parse_list_string",
+]
+
+
+def parse_list_string(val: str) -> list | None:
+    if pd.isna(val):
+        return None
+    try:
+        parsed = ast.literal_eval(val)
+        if isinstance(parsed, list):
+            return parsed
+        return [parsed]
+    except (ValueError, SyntaxError):
+        return None
+
+
+def parse_first_element(val: str) -> float:
+    if pd.isna(val):
+        return float("nan")
+    try:
+        parsed = ast.literal_eval(val)
+        if isinstance(parsed, list):
+            if len(parsed) > 0:
+                return float(parsed[0])
+            return float("nan")
+        return float(parsed)
+    except (ValueError, SyntaxError):
+        return float("nan")
+
+
+def sum_list_elements(val: str) -> float:
+    if pd.isna(val):
+        return float("nan")
+    try:
+        parsed = ast.literal_eval(val)
+        if isinstance(parsed, list):
+            if len(parsed) > 0:
+                try:
+                    return float(sum(float(item) for item in parsed))
+                except (ValueError, TypeError):
+                    return float("nan")
+            return float("nan")
+        return float(parsed)
+    except (ValueError, SyntaxError):
+        return float("nan")
+
+
+def is_homogeneous(val: str) -> bool:
+    if pd.isna(val):
+        return False
+    try:
+        parsed = ast.literal_eval(val)
+        if isinstance(parsed, list):
+            if len(parsed) > 0:
+                return len(set(parsed)) == 1
+            return False
+        return True
+    except (ValueError, SyntaxError, TypeError):
+        return False

dr_frames/profiling.py

@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Callable
+from pathlib import PurePath
+from typing import Any
+
+import pandas as pd
+from pydantic import BaseModel, Field
+
+from .aggregation import unique_non_null
+
+__all__ = [
+    "ColInfo",
+    "DFColInfo",
+    "looks_like_json",
+    "looks_like_path",
+    "infer_series_base_tag_type",
+    "infer_tags_from_series_sample",
+    "infer_col_name_contains_tags",
+    "infer_col_name_suffix_tags",
+    "infer_col_name_prefix_tags",
+]
+
+
+class ColInfo(BaseModel):
+    catalog: DFColInfo
+    name: str
+    dtype: str | None = None
+    tags: set[str] = Field(default_factory=set)
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    def update_tags(self, series: pd.Series) -> None:
+        self.tags.update(
+            infer_series_base_tag_type(
+                series,
+                self.catalog.pd_type_to_tags,
+                self.catalog.pd_default_type,
+            )
+        )
+        self.tags.update(
+            infer_col_name_contains_tags(
+                self.name,
+                self.catalog.col_name_contains_map,
+            )
+        )
+        self.tags.update(
+            infer_col_name_suffix_tags(
+                self.name,
+                self.catalog.col_name_suffix_map,
+            )
+        )
+        self.tags.update(
+            infer_col_name_prefix_tags(
+                self.name,
+                self.catalog.col_name_prefix_map,
+            )
+        )
+        self.tags.update(
+            infer_tags_from_series_sample(
+                series,
+                self.catalog.path_like_extensions,
+            )
+        )
+
+    def has_tag(self, tag: str) -> bool:
+        return tag in self.tags
+
+    def add_tags(self, extra: list[str]) -> None:
+        self.tags.update(extra)
+
+
+class DFColInfo(BaseModel):
+    columns: dict[str, ColInfo] = Field(default_factory=dict)
+    pd_type_to_tags: list[tuple[Callable, list[str]]] = Field(
+        default_factory=lambda: [
+            (pd.api.types.is_bool_dtype, ["bool"]),
+            (pd.api.types.is_integer_dtype, ["int", "numeric"]),
+            (pd.api.types.is_float_dtype, ["float", "numeric"]),
+            (pd.api.types.is_datetime64_any_dtype, ["datetime"]),
+            (pd.api.types.is_timedelta64_dtype, ["timedelta"]),
+            (lambda dt: isinstance(dt, pd.CategoricalDtype), ["categorical"]),
+        ]
+    )
+    pd_default_type: list[str] = Field(default_factory=lambda: ["str"])
+
+    col_name_contains_map: dict[tuple[str, ...], str] = Field(
+        default_factory=lambda: {
+            ("config", "kwargs", "settings", "params"): "config",
+        }
+    )
+    col_name_suffix_map: dict[tuple[str, ...], str] = Field(
+        default_factory=lambda: {
+            ("_path", "_dir"): "path",
+            ("_id",): "id",
+            ("_tag", "_tags"): "categorical",
+        }
+    )
+    col_name_prefix_map: dict[tuple[str, ...], str] = Field(
+        default_factory=lambda: {
+            ("is_", "has_"): "bool_like",
+            ("metric_",): "metric",
+        }
+    )
+
+    path_like_extensions: set[str] = Field(
+        default_factory=lambda: {
+            ".json",
+            ".jsonl",
+            ".csv",
+            ".tsv",
+            ".parquet",
+            ".txt",
+            ".yaml",
+            ".yml",
+            ".log",
+        }
+    )
+
+    def update_from_df(self, df: pd.DataFrame) -> None:
+        for column in df.columns:
+            info = ColInfo(catalog=self, name=column)
+            info.update_tags(df[column])
+            self.columns[column] = info
+
+    def get(self, name: str) -> ColInfo | None:
+        return self.columns.get(name)
+
+    def names_with_tag(self, tag: str) -> list[str]:
+        return [name for name, info in self.columns.items() if info.has_tag(tag)]
+
+
+def infer_tags_from_series_sample(
+    series: pd.Series,
+    path_like_extensions: set[str],
+    sample_size: int = 10,
+) -> set[str]:
+    tags: set[str] = set()
+    unique_vals = unique_non_null(series)
+    if len(unique_vals) == 0:
+        return tags
+
+    # Sample up to sample_size unique non-null values
+    sample_vals = unique_vals[: min(sample_size, len(unique_vals))]
+
+    # Check each sampled value for path and JSON patterns
+    path_matches = 0
+    json_matches = 0
+    for value in sample_vals:
+        if looks_like_path(value, path_like_extensions):
+            path_matches += 1
+        if looks_like_json(value):
+            json_matches += 1
+
+    # Add tags if at least one value matches
+    if path_matches > 0:
+        tags.add("path")
+    if json_matches > 0:
+        tags.add("json")
+    return tags
+
+
+def infer_col_name_contains_tags(
+    name: str,
+    col_name_contains_map: dict[tuple[str, ...], str],
+) -> set[str]:
+    lower_name = name.lower()
+    tags: set[str] = set()
+    for contains_tuple, tag in col_name_contains_map.items():
+        for contains in contains_tuple:
+            if contains.lower() in lower_name:
+                tags.add(tag)
+    return tags
+
+
+def infer_col_name_suffix_tags(
+    name: str,
+    col_name_suffix_map: dict[tuple[str, ...], str],
+) -> set[str]:
+    lower_name = name.lower()
+    tags: set[str] = set()
+    for suffix_tuple, tag in col_name_suffix_map.items():
+        for suffix in suffix_tuple:
+            if lower_name.endswith(suffix.lower()):
+                tags.add(tag)
+                break
+    return tags
+
+
+def infer_col_name_prefix_tags(
+    name: str,
+    col_name_prefix_map: dict[tuple[str, ...], str],
+) -> set[str]:
+    lower_name = name.lower()
+    tags: set[str] = set()
+    for prefix_tuple, tag in col_name_prefix_map.items():
+        for prefix in prefix_tuple:
+            if lower_name.startswith(prefix.lower()):
+                tags.add(tag)
+                break
+    return tags
+
+
+def infer_series_base_tag_type(
+    series: pd.Series,
+    pd_type_to_tags: list[tuple[Callable, list[str]]],
+    pd_default_type: list[str],
+) -> set[str]:
+    dtype = series.dtype
+    tags: set[str] = set()
+    for type_fxn, tag_list in pd_type_to_tags:
+        tags.update(tag_list if type_fxn(dtype) else [])
+    if len(tags) == 0:
+        tags.update(pd_default_type)
+    if series.isna().any():
+        tags.add("nullable")
+    return tags
+
+
+def looks_like_path(value: Any, path_extensions: set[str]) -> bool:
+    if not value or not isinstance(value, str) or value.isspace():
+        return False
+    normalized = value.strip()
+    if any(sep in normalized for sep in ("/", "\\")):
+        return True
+    try:
+        suffix = PurePath(normalized).suffix.lower()
+    except (TypeError, ValueError):
+        return False
+    return suffix in path_extensions
+
+
+def looks_like_json(value: Any) -> bool:
+    if not value or not isinstance(value, str):
+        return False
+    trimmed = value.strip()
+    if not trimmed:
+        return False
+    if not (trimmed.startswith("{") or trimmed.startswith("[")):
+        return False
+    try:
+        parsed = json.loads(trimmed)
+    except Exception:
+        return False
+    return isinstance(parsed, (dict, list))