diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index 39690aa0e..0248630f2 100644 --- a/.github/workflows/checks.yml +++ b/.github/workflows/checks.yml @@ -34,6 +34,41 @@ jobs: - name: Run paddy-format check run: find buckaroo tests scripts -type f -name "*.py" -not -path "buckaroo/jlisp/lispy.py" -print0 | xargs -0 uv run python scripts/paddy_format.py --check + # Non-blocking static type check (continue-on-error) over the files touched + # by the dataflow generic-typing work. Scoped + standard mode are both + # deliberate: pyrightconfig.typecheck.json lists exactly those files and runs + # basedpyright's "standard" mode rather than its stricter "recommended" + # default, so the signal stays on real type errors instead of third-party + # stub noise. Editors keep using the [tool.basedpyright] block in + # pyproject.toml — that file's name is not auto-discovered. + TypeCheck: + name: Python / Typecheck (non-blocking) + runs-on: depot-ubuntu-latest + timeout-minutes: 5 + continue-on-error: true + steps: + - uses: actions/checkout@v6 + - name: Install uv + uses: astral-sh/setup-uv@v7 + with: + enable-cache: true + - name: Install the project + # --all-extras so basedpyright can resolve polars/xorq/pyarrow imports. + run: uv sync --locked --all-extras --dev + - name: Run basedpyright (scoped, informational) + # Informational only: these files carry a large pre-existing backlog of + # standard-mode diagnostics, so the step is forced green and the totals + # go to the job summary. The full report is in the step log. `|| true` + # keeps the check green; job-level continue-on-error covers setup steps. + run: | + uv run --with basedpyright==1.39.8 basedpyright --project pyrightconfig.typecheck.json 2>&1 | tee bpr.log || true + { + echo "## basedpyright (standard mode, dataflow typing files — non-blocking)" + echo '```' + tail -n 1 bpr.log + echo '```' + } >> "$GITHUB_STEP_SUMMARY" + TestJS: name: JS / Build + Test runs-on: depot-ubuntu-latest diff --git a/buckaroo/buckaroo_widget.py b/buckaroo/buckaroo_widget.py index d6f3d0406..346e765c7 100644 --- a/buckaroo/buckaroo_widget.py +++ b/buckaroo/buckaroo_widget.py @@ -11,7 +11,7 @@ import sys import traceback from datetime import datetime -from typing import Literal, Union +from typing import Any as TAny, ClassVar, Literal, Type, Union import pandas as pd import json import logging @@ -35,6 +35,10 @@ from .dataflow.autocleaning import PandasAutocleaning from pathlib import Path +# pandas backend binding of the generic dataflow. Mirrors PolarsDataflow +# (polars_buckaroo) and XorqDataflow (xorq_buckaroo). +PandasDataflow = CustomizableDataflow[pd.DataFrame] + logger = logging.getLogger() # Opt-in cross-boundary tracing. Set BUCKAROO_BK_FLASH=1 in the kernel @@ -184,7 +188,10 @@ def _df_to_obj(self, df:pd.DataFrame): autocleaning_klass = PandasAutocleaning #override the base CustomizableDataFlow klass DFStatsClass = DfStatsV2 # Pandas Specific autoclean_conf = tuple([CleaningConf, NoCleaningConf]) #override the base CustomizableDataFlow conf - dataflow_klass = CustomizableDataflow + # The composed dataflow's backend binding. Typed broadly here (one slot, + # specialised per backend subclass) — the precise binding lives in the + # named alias/subclass assigned: PandasDataflow / PolarsDataflow / XorqDataflow. + dataflow_klass: ClassVar[Type[CustomizableDataflow[TAny]]] = PandasDataflow df_data_dict = Dict({}).tag(sync=True) diff --git a/buckaroo/dataflow/abc_dataflow.py b/buckaroo/dataflow/abc_dataflow.py index c0e18b3e5..d9584b84f 100644 --- a/buckaroo/dataflow/abc_dataflow.py +++ b/buckaroo/dataflow/abc_dataflow.py @@ -3,41 +3,72 @@ from __future__ import annotations from abc import ABCMeta, abstractmethod -from typing import Any, Dict, List, Type +from typing import Any, Generic, List, Optional, Type from traitlets import HasTraits, MetaHasTraits from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis +from .df_types import FrameT + class _ABCMetaHasTraits(ABCMeta, MetaHasTraits): pass -class ABCDataflow(HasTraits, metaclass=_ABCMetaHasTraits): +class ABCDataflow(HasTraits, Generic[FrameT], metaclass=_ABCMetaHasTraits): """ Abstract base for dataflow implementations. Reference implementations: DataFlow and CustomizableDataflow. Other implementations (e.g., ColumnExecutorDataflow) should conform. + + Generic on ``FrameT`` — the (unbounded) carrier type, so this umbrella + covers both the eager subtree (``DataFlow`` / ``CustomizableDataflow``, + which narrow to the ``DataFrameLike``-bound ``DataFrameT``) and the lazy + ``ColumnExecutorDataflow`` (``ABCDataflow[pl.LazyFrame]``). See + ``df_types`` for the ``FrameT`` / ``DataFrameT`` split. """ # Baseline interface expected by Buckaroo widgets and extensions. + # + # The synced/computed wire attributes below are declared ``Any`` on + # purpose. Every concrete dataflow backs them with a traitlets trait + # (``Dict(...)`` / ``Any(...)``) or a ``@property``, and a traitlets + # descriptor returns ``Any`` on instance access. Declaring a concrete + # ``Dict[str, Any]`` here only fights the descriptor protocol: the + # subclass trait/property reads as an incompatible override, and a write + # (``self.summary_sd = ...``) reads as an illegal assignment. What these + # document is "this name exists on every dataflow"; each trait's actual + # shape lives with its definition. Same rationale as the frame-carrying + # methods in ``df_types`` — type where it's sound, not on the traits. analysis_klasses: List[Type[ColAnalysis]] - df_data_dict: Dict[str, Any] - df_display_args: Dict[str, Any] - df_meta: Dict[str, Any] - buckaroo_options: Dict[str, Any] - command_config: Dict[str, Any] + df_data_dict: Any + df_display_args: Any + df_meta: Any + buckaroo_options: Any + command_config: Any operations: Any - operation_results: Dict[str, Any] - summary_sd: Dict[str, Any] - cleaned_sd: Dict[str, Any] - processed_sd: Dict[str, Any] - merged_sd: Dict[str, Any] + operation_results: Any + summary_sd: Any + cleaned_sd: Any + processed_sd: Any + merged_sd: Any widget_args_tuple: Any - processed_df: Any - - analysis_klasses: List[Type[ColAnalysis]] + + @property + @abstractmethod + def processed_df(self) -> Optional[FrameT]: + """The fully-processed frame the widget renders, in the backend's + own type (``None`` before the pipeline has run, or for lazy + backends that never materialize). + + A read-only computed property in every implementation — eager + backends derive it from ``processed_result``; the lazy + ``ColumnExecutorDataflow`` never materializes and returns ``None``. + Declared as a property (not a plain attribute) so those + ``@property`` overrides are a compatible, like-for-like override. + """ + ... @abstractmethod def populate_df_meta(self) -> None: diff --git a/buckaroo/dataflow/column_executor_dataflow.py b/buckaroo/dataflow/column_executor_dataflow.py index a1290de9f..e0b58f9a9 100644 --- a/buckaroo/dataflow/column_executor_dataflow.py +++ b/buckaroo/dataflow/column_executor_dataflow.py @@ -27,9 +27,15 @@ from buckaroo.file_cache.batch_planning import PlanningFunction -class ColumnExecutorDataflow(ABCDataflow): +class ColumnExecutorDataflow(ABCDataflow[pl.LazyFrame]): """A minimal DataFlow focused on column-executor-driven summary stats for Polars LazyFrames. + Binds the abstract base's unbounded ``FrameT`` to ``pl.LazyFrame``. A + LazyFrame is never materialised, so it cannot meet the eager + ``DataFrameLike`` contract (no ``len`` / row-slice) — which is why this + class inherits ``ABCDataflow`` directly rather than the eager + ``CustomizableDataflow`` body, and supplies its own executor pipeline. + - Works with a LazyFrame and avoids materializing the dataframe on load. - No-op command config, autocleaning, quick commands, and @@ -364,7 +370,9 @@ def auto_compute_summary(self, sync_executor_class: Type[Executor], parallel_exe # Don't re-raise, let it fail silently and use defaults @property - def processed_df(self) -> Any: + def processed_df(self) -> Optional[pl.LazyFrame]: + # Lazy: the frame is never materialised, so there is no processed + # frame to render. Always None (matches ABCDataflow's Optional[FrameT]). return None @observe('merged_sd') diff --git a/buckaroo/dataflow/dataflow.py b/buckaroo/dataflow/dataflow.py index 3d8b33047..4dcb5a681 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -1,4 +1,6 @@ -from typing import List, Literal, Tuple, Type, TypedDict, Dict as TDict, Any as TAny, Union +from typing import ( + Generic, List, Literal, Optional, Protocol, Tuple, Type, TypedDict, cast, + Dict as TDict, Any as TAny, Union) from typing_extensions import override import six import warnings @@ -21,6 +23,7 @@ from .abc_dataflow import ABCDataflow +from .df_types import DataFrameT from .sd_cache import hash_chain, split_chain_by_scope @@ -41,15 +44,38 @@ class DfTrait(Any): def set(self, obj, value): new_value = self._validate(obj, value) + # self.name is Optional[str] in the traitlets stubs (None until the + # metaclass binds the trait to a class), but it is always set by the + # time set() runs on an instance. + name = cast(str, self.name) try: - old_value = obj._trait_values[self.name] + old_value = obj._trait_values[name] except KeyError: old_value = self.default_value - obj._trait_values[self.name] = new_value + obj._trait_values[name] = new_value if old_value is not new_value: - obj._notify_trait(self.name, old_value, new_value) + obj._notify_trait(name, old_value, new_value) -class DataFlow(ABCDataflow): +class Autocleaning(Protocol): + """Structural interface the dataflow pipeline calls on ``self.ac_obj``. + + The eager backends (``PandasAutocleaning`` and the polars/xorq variants) + supply all of it. The bare-pipeline ``SentinelAutocleaning`` implements + only the subset the stub path touches (``command_config`` / + ``handle_ops_and_clean``), so ``ac_obj`` is cast to this Protocol at + construction — the code-interpreter members below are reached only via + ``CustomizableDataflow``, which always runs on a real backend. + """ + command_config: TAny + config_dict: TAny + def handle_ops_and_clean(self, df: TAny, cleaning_method: TAny, + quick_command_args: TAny, existing_operations: TAny) -> TAny: ... + def add_command(self, incomingCommandKls: TAny) -> TAny: ... + def _run_df_interpreter(self, df: TAny, operations: TAny, initial_sd: TAny) -> TAny: ... + def run_code_generator(self, operations: TAny) -> TAny: ... + + +class DataFlow(ABCDataflow[DataFrameT], Generic[DataFrameT]): """This class is meant to only represent the dataflow through buckaroo with no accomodation for widget particulars @@ -67,15 +93,24 @@ class DataFlow(ABCDataflow): """ def __init__(self, raw_df): - self.exception = None + # Set by exception_protect handlers (see dataflow_extras) to + # sys.exc_info() when a trait-observer cascade fails; None until then. + self.exception: Optional[Tuple[TAny, TAny, TAny]] = None super().__init__() self.summary_sd = {} self.existing_operations = [] - self.ac_obj = self.autocleaning_klass(self.autoclean_conf) + # SentinelAutocleaning (the bare default) implements only part of the + # Autocleaning interface; the real backends supply all of it. Cast so + # the code-interpreter methods type-check on the customizable path. + self.ac_obj = cast(Autocleaning, self.autocleaning_klass(self.autoclean_conf)) self.command_config = self.ac_obj.command_config try: self.raw_df = raw_df except Exception: + if self.exception is None: + # No handler stored the original exc_info — propagate the + # current exception rather than subscripting None. + raise six.reraise(self.exception[0], self.exception[1], self.exception[2]) autocleaning_klass = SentinelAutocleaning @@ -102,7 +137,12 @@ def __init__(self, raw_df): post_processing_method = Unicode('').tag(default='') processed_result = DfTrait().tag(default=None) - analysis_klasses = None + # Bare-pipeline stub: no analyses by default. ``[]`` (not ``None``) so the + # type stays ``List[Type[ColAnalysis]]`` end-to-end — CustomizableDataflow + # narrows it to a real list, and nothing relies on a None sentinel (reads + # are id()-based cache keys and the dead "foo"/"bar" test branches below, + # all of which treat None and [] alike). + analysis_klasses: List[Type[ColAnalysis]] = [] summary_sd = Any() df_meta = Any() @@ -133,7 +173,7 @@ def __init__(self, raw_df): - def _compute_sampled_df(self, raw_df:pd.DataFrame, sample_method:str): + def _compute_sampled_df(self, raw_df: DataFrameT, sample_method: str) -> DataFrameT: if sample_method == "first": return raw_df[:1] return raw_df @@ -192,9 +232,13 @@ def _operation_result(self, _change:Any) -> None: self._in_operation_result = False @property - def cleaned_df(self): + def cleaned_df(self) -> Optional[DataFrameT]: + # traitlets descriptors return Any on access, so the element type + # of the ``cleaned`` tuple is opaque to the checker — cast to the + # frame type this dataflow is bound to. if self.cleaned is not None: - return self.cleaned[0] + return cast(DataFrameT, self.cleaned[0]) + return None @property def cleaned_sd(self): @@ -212,8 +256,10 @@ def merged_operations(self): if self.cleaned is not None: return self.cleaned[3] - def _compute_processed_result(self, cleaned_df, post_processing_method): - return [cleaned_df, {}] + def _compute_processed_result( + self, cleaned_df: DataFrameT, + post_processing_method: str) -> Tuple[DataFrameT, SDType]: + return (cleaned_df, {}) def populate_df_meta(self): pass @@ -226,9 +272,9 @@ def _processed_result(self, change): self.populate_df_meta() @property - def processed_df(self): + def processed_df(self) -> Optional[DataFrameT]: if self.processed_result is not None: - return self.processed_result[0] + return cast(DataFrameT, self.processed_result[0]) return None @property @@ -237,14 +283,14 @@ def processed_sd(self) -> SDType: return self.processed_result[1] return {} - def _get_summary_sd(self, df:pd.DataFrame) -> Tuple[SDType, TAny]: + def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, TAny]: analysis_klasses = self.analysis_klasses if analysis_klasses == "foo": return {'some-col': {'foo':8}}, {} if analysis_klasses == "bar": return {'other-col': {'bar':10}}, {} ret_summary = {} - for col in df.columns: + for col in processed_df.columns: ret_summary[col] = {} return ret_summary, {} @@ -260,6 +306,11 @@ def _summary_sd(self, change): # Skip when neither the dataframe nor analysis_klasses has actually # changed since the last run. See issue #709. df = self.processed_df + if df is None: + # Nothing to summarize before the pipeline has produced a frame. + # Matches the guards in _merged_sd and _populate_sd_cache, and + # narrows Optional[DataFrameT] -> DataFrameT for _get_summary_sd. + return klasses = self.analysis_klasses if (id(df), id(klasses)) == self._summary_sd_cache_key: return @@ -298,9 +349,14 @@ def _widget_config(self, change): 'summary_stats': List[str]}) -class CustomizableDataflow(DataFlow): +class CustomizableDataflow(DataFlow[DataFrameT], Generic[DataFrameT]): """ This allows targetd extension and customization of DataFlow + + Still generic on ``DataFrameT``: both the pandas and polars backends + use this class directly (bound to ``pd.DataFrame`` / ``pl.DataFrame`` + respectively), and ``XorqDataflow`` subclasses it as + ``CustomizableDataflow[XorqExpr]``. """ #analysis_klasses = [StylingAnalysis] analysis_klasses: List[Type[ColAnalysis]] = [StylingAnalysis] @@ -372,9 +428,11 @@ def populate_df_meta(self) -> None: 'rows_shown': min(len(self.processed_df), self.sampling_klass.serialize_limit), 'total_rows': len(self.orig_df)} - #typing compalins about this, but so far as this class is concerned, buckaroo_options follows theBuckarooOptions type - # typing doesn't get along well with traitlets - buckaroo_options:BuckarooOptions = Dict({'sampled': ['random'], 'auto_clean': ['aggressive', 'conservative'], + # buckaroo_options is BuckarooOptions-shaped at runtime, but it's a + # traitlets ``Dict`` trait, so we let it inherit the base ``Any`` rather + # than annotate ``: BuckarooOptions`` — the descriptor value can't + # satisfy that declared type. See ABCDataflow's wire-attribute note. + buckaroo_options = Dict({'sampled': ['random'], 'auto_clean': ['aggressive', 'conservative'], 'post_processing': [], 'df_display': ['main', 'summary'], 'show_commands': ['on'], 'summary_stats': ['all']}).tag(sync=True) def setup_options_from_analysis(self): @@ -609,7 +667,9 @@ def run_code_generator(self, operations): ### end code interpeter block @override - def _compute_processed_result(self, cleaned_df:pd.DataFrame, post_processing_method:str) -> Tuple[pd.DataFrame, SDType]: + def _compute_processed_result( + self, cleaned_df: DataFrameT, + post_processing_method: str) -> Tuple[DataFrameT, SDType]: if post_processing_method == '': return (cleaned_df, {}) else: @@ -618,16 +678,19 @@ def _compute_processed_result(self, cleaned_df:pd.DataFrame, post_processing_met ret_df, sd = post_analysis.post_process_df(cleaned_df) return (ret_df, sd) except Exception as e: - return (self._build_error_dataframe(e), {}) + # Error frames are always pandas regardless of backend — the + # styling layer renders them through the pandas path. The cast + # documents that intentional impurity in the DataFrameT contract. + return (cast(DataFrameT, self._build_error_dataframe(e)), {}) - def _build_error_dataframe(self, e): + def _build_error_dataframe(self, e) -> pd.DataFrame: return pd.DataFrame({'err': [str(e)]}) ### start summary stats block #TAny closer to some error type @override - def _get_summary_sd(self, processed_df:pd.DataFrame) -> Tuple[SDType, TDict[str, TAny]]: + def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, TDict[str, TAny]]: stats = self.DFStatsClass( processed_df, self.analysis_klasses, @@ -659,7 +722,7 @@ def _sd_to_jsondf(self, sd:SDType): keep = wire_stat_keys(self.df_display_klasses.values(), self.pinned_rows) return sd_to_parquet_b64(project_sd(sd, keep)) - def _df_to_obj(self, df:pd.DataFrame) -> TDict[str, TAny]: + def _df_to_obj(self, df: DataFrameT) -> TDict[str, TAny]: return pd_to_obj(self.sampling_klass.serialize_sample(df)) def add_analysis(self, analysis_klass:Type[ColAnalysis]) -> None: diff --git a/buckaroo/dataflow/df_types.py b/buckaroo/dataflow/df_types.py new file mode 100644 index 000000000..b652a7fcd --- /dev/null +++ b/buckaroo/dataflow/df_types.py @@ -0,0 +1,92 @@ +"""Backend-agnostic dataframe typing for the dataflow pipeline. + +The dataflow classes move a single "dataframe" object through a fixed +pipeline (``raw_df`` -> ``sampled_df`` -> ``cleaned`` -> ``processed``) +and then derive summary statistics from it. Three concrete backends +supply that object, and they share *no* common base class: + +- pandas -> ``pandas.DataFrame`` +- polars -> ``polars.DataFrame`` +- xorq -> a xorq/ibis expression (``xorq.vendor.ibis`` ``Table``) + +There are actually two tiers of carrier, so there are two type variables: + +``FrameT`` (unbounded) — the umbrella for *every* dataflow, including the +lazy one. ``ColumnExecutorDataflow`` carries a ``pl.LazyFrame`` that is +never materialised: it has no row count and is not row-sliceable, so it +cannot meet the eager contract below. It binds ``FrameT = pl.LazyFrame`` +directly on the abstract base and supplies its own pipeline. + +``DataFrameT`` (bound to ``DataFrameLike``) — the *eager*, materialised +frames used by the shared ``DataFlow`` / ``CustomizableDataflow`` body, +which reads ``df.columns``, calls ``len(df)``, and row-slices ``df[:n]``. +pandas, polars (eager), and xorq ``Table`` all satisfy this structurally. +Each eager backend binds it to its concrete frame type:: + + CustomizableDataflow[pd.DataFrame] # pandas + CustomizableDataflow[pl.DataFrame] # polars (eager) + XorqDataflow == CustomizableDataflow[XorqExpr] # xorq + +There is no nominal base class spanning these three, so the bound is a +structural ``Protocol``, not a shared superclass. ``DataFrameT`` is a +valid argument to ``FrameT`` (a bounded type var satisfies an unbounded +one), so ``DataFlow(ABCDataflow[DataFrameT])`` type-checks while the lazy +``ColumnExecutorDataflow(ABCDataflow[pl.LazyFrame])`` also does. + +The type variable flows through the method boundaries that actually take +or return the frame — ``_compute_sampled_df``, ``_compute_processed_result``, +``_get_summary_sd``, and the ``processed_df`` / ``cleaned_df`` properties +— so a ``CustomizableDataflow[pd.DataFrame]`` reports ``processed_df`` as +``pd.DataFrame | None`` rather than ``Any``. + +The traitlets-backed attributes (``raw_df``, ``sampled_df``, ``cleaned``, +``processed_result``, the lazy ``raw_ldf``) stay deliberately untyped: a +traitlets descriptor returns ``Any`` on instance access, and annotating +the class-level trait assignment fights the descriptor protocol. The +generic contract is expressed where it is sound — the method boundaries — +not on the traits. +""" +from __future__ import annotations + +from typing import Any, Protocol, TypeVar + + +class DataFrameLike(Protocol): + """The structural surface the *shared* dataflow body uses on a frame. + + The pandas/polars-shared code in ``CustomizableDataflow`` reads + ``df.columns``, calls ``len(df)``, and row-slices ``df[:n]``. All three + backends expose these members, so each backend's concrete type + structurally satisfies this Protocol and can bind ``DataFrameT``. + + Caveat for xorq: an ibis/xorq expression *defines* ``__len__`` but + raises at runtime (it is not materialised, so it has no row count). + That is fine for the static bound — and ``XorqDataflow`` overrides the + methods that would actually call ``len`` (``populate_df_meta`` issues + ``expr.count().execute()`` instead). The Protocol describes the + interface the shared code *names*; backends whose semantics differ + override those methods rather than weakening the bound. + """ + + @property + def columns(self) -> Any: ... + def __len__(self) -> int: ... + def __getitem__(self, key: Any) -> Any: ... + + +#: Unbounded umbrella carrier for the abstract base ``ABCDataflow`` — spans +#: every dataflow, including the lazy ``ColumnExecutorDataflow`` whose +#: ``pl.LazyFrame`` cannot meet the eager ``DataFrameLike`` contract (it is +#: never materialised). Eager subclasses pass the tighter ``DataFrameT`` in +#: its place; the lazy one binds ``pl.LazyFrame`` directly. +FrameT = TypeVar("FrameT") + +#: The *eager* dataframe type carried through the shared ``DataFlow`` / +#: ``CustomizableDataflow`` body. Bound to ``DataFrameLike`` — the minimal +#: structural interface that body uses (``columns`` / ``len`` / row-slice) +#: — which pandas frames, polars (eager) frames, and xorq ``Table`` +#: expressions all satisfy. There is no nominal base class spanning the +#: three (see module docstring), so the bound is structural. +DataFrameT = TypeVar("DataFrameT", bound=DataFrameLike) + +__all__ = ["FrameT", "DataFrameT", "DataFrameLike"] diff --git a/buckaroo/polars_buckaroo.py b/buckaroo/polars_buckaroo.py index 0742afc0b..e9640f79a 100644 --- a/buckaroo/polars_buckaroo.py +++ b/buckaroo/polars_buckaroo.py @@ -11,11 +11,15 @@ from .serialization_utils import pd_to_obj from .customizations.styling import DefaultSummaryStatsStyling, DefaultMainStyling from .customizations.pl_autocleaning_conf import NoCleaningConfPl -from .dataflow.dataflow import Sampling +from .dataflow.dataflow import CustomizableDataflow, Sampling from .dataflow.autocleaning import PandasAutocleaning from .dataflow.widget_extension_utils import configure_buckaroo from .styling_helpers import obj_, pinned_histogram, pinned_filtered_histogram +# polars backend binding of the generic dataflow. Mirrors PandasDataflow +# (buckaroo_widget) and XorqDataflow (xorq_buckaroo). +PolarsDataflow = CustomizableDataflow[pl.DataFrame] + class PLSampling(Sampling): pre_limit = False serialize_limit = 1_000_000 @@ -73,6 +77,7 @@ class PolarsBuckarooWidget(BuckarooWidget): autoclean_conf = tuple([NoCleaningConfPl]) #override the base CustomizableDataFlow conf DFStatsClass = PlDfStatsV2 sampling_klass = PLSampling + dataflow_klass = PolarsDataflow # composed dataflow carries polars DataFrames # _sd_to_jsondf is inherited from BuckarooWidgetBase, which delegates to # the dataflow so the wire-stat projection (#880) lives in one place. diff --git a/buckaroo/server/data_loading.py b/buckaroo/server/data_loading.py index 9636e7f3b..12a890d9f 100644 --- a/buckaroo/server/data_loading.py +++ b/buckaroo/server/data_loading.py @@ -36,7 +36,7 @@ def pre_stats_sample(kls, df): return df -class ServerDataflow(CustomizableDataflow): +class ServerDataflow(CustomizableDataflow[pd.DataFrame]): """Headless dataflow matching BuckarooInfiniteWidget's pipeline.""" sampling_klass = ServerSampling autocleaning_klass = PandasAutocleaning diff --git a/buckaroo/server/data_loading_polars.py b/buckaroo/server/data_loading_polars.py index 26f60df6e..c3de0d5b2 100644 --- a/buckaroo/server/data_loading_polars.py +++ b/buckaroo/server/data_loading_polars.py @@ -39,7 +39,7 @@ class PolarsServerSampling(PLSampling): serialize_limit = -1 # infinite mode — no per-page sample cap -class PolarsServerDataflow(CustomizableDataflow): +class PolarsServerDataflow(CustomizableDataflow[pl.DataFrame]): """Headless polars dataflow matching ``PolarsBuckarooInfiniteWidget``.""" analysis_klasses = local_analysis_klasses autocleaning_klass = PandasAutocleaning diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index 2ffd74305..1b3430042 100644 --- a/buckaroo/xorq_buckaroo.py +++ b/buckaroo/xorq_buckaroo.py @@ -14,9 +14,12 @@ import traceback import weakref from io import BytesIO -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd + +if TYPE_CHECKING: + from xorq.vendor.ibis.expr.types.relations import Table as XorqExpr import pyarrow as pa import pyarrow.parquet as pq from traitlets import Unicode @@ -129,7 +132,7 @@ class XorqAutocleaning(PandasAutocleaning): """ -class XorqDataflow(CustomizableDataflow): +class XorqDataflow(CustomizableDataflow["XorqExpr"]): """Dataflow specialised for ibis/xorq expression inputs. Two pieces of behaviour differ from the pandas dataflow: @@ -156,7 +159,7 @@ def populate_df_meta(self) -> None: 'rows_shown': rows_shown, 'total_rows': _expr_count(self.orig_df)} - def _get_summary_sd(self, processed_df): + def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): if _is_pandas(processed_df): # The error path (and any postprocessor that returns a pandas # DataFrame) doesn't go through XorqStatPipeline. Return a diff --git a/pyrightconfig.typecheck.json b/pyrightconfig.typecheck.json new file mode 100644 index 000000000..6c67d6412 --- /dev/null +++ b/pyrightconfig.typecheck.json @@ -0,0 +1,17 @@ +{ + "typeCheckingMode": "standard", + "pythonVersion": "3.11", + "reportMissingImports": "error", + "reportMissingTypeStubs": false, + "include": [ + "buckaroo/dataflow/df_types.py", + "buckaroo/dataflow/abc_dataflow.py", + "buckaroo/dataflow/dataflow.py", + "buckaroo/dataflow/column_executor_dataflow.py", + "buckaroo/buckaroo_widget.py", + "buckaroo/polars_buckaroo.py", + "buckaroo/xorq_buckaroo.py", + "buckaroo/server/data_loading.py", + "buckaroo/server/data_loading_polars.py" + ] +}