preview_results.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from data_designer.config.analysis.dataset_profiler import DatasetProfilerResults
from data_designer.config.config_builder import DataDesignerConfigBuilder
from data_designer.config.dataset_metadata import DatasetMetadata
from data_designer.config.utils.visualization import WithRecordSamplerMixin

if TYPE_CHECKING:
    import pandas as pd


class PreviewResults(WithRecordSamplerMixin):
    def __init__(
        self,
        *,
        config_builder: DataDesignerConfigBuilder,
        dataset_metadata: DatasetMetadata | None = None,
        dataset: pd.DataFrame | None = None,
        analysis: DatasetProfilerResults | None = None,
        processor_artifacts: dict[str, list[dict]] | None = None,
        task_traces: list[Any] | None = None,
    ):
        """Creates a new instance with results from a Data Designer preview run.

        Args:
            config_builder: Data Designer configuration builder.
            dataset_metadata: Metadata about the generated dataset (e.g., seed column names).
            dataset: Dataset of the preview run.
            analysis: Analysis of the preview run.
            processor_artifacts: Artifacts generated by the processors.
            task_traces: Async scheduler task traces (when DATA_DESIGNER_ASYNC_TRACE=1).
        """
        self.dataset: pd.DataFrame | None = dataset
        self.analysis: DatasetProfilerResults | None = analysis
        self.processor_artifacts: dict[str, list[dict]] | None = processor_artifacts
        self.dataset_metadata: DatasetMetadata | None = dataset_metadata
        self.task_traces: list[Any] | None = task_traces
        self._config_builder = config_builder