Source code for bukka.coding.write_pipeline

from typing import Any
from pathlib import Path
from bukka.coding.utils.template_handler import TemplateBaseClass

# Template for the complete pipeline file
FULL_TEMPLATE = '''
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
{imports}

{instantiations}

preprocessor = ColumnTransformer([
    {transformer_steps}
])

pipeline = Pipeline([
    {pipeline}
])
'''

SEPARATOR = ",\n\t"


[docs] class PipelineWriter(TemplateBaseClass): """Build a pipeline plan from identified problems and chosen solutions. The writer takes pipeline steps and generates a complete sklearn Pipeline with appropriate imports, instantiations, ColumnTransformer (if needed), and final Pipeline definition. Parameters ---------- pipeline_steps : list[tuple[Any, Any]] List of tuples where each tuple contains (solution_object, problem_object). Solution objects must have `fetch_import()`, `fetch_instantiation()`, and `name` attributes. Problem objects must have `problem_type` and `features` attributes. output_path : str | Path The file path where the generated pipeline code will be written. Attributes ---------- pipeline_steps : list[tuple[Any, Any]] The pipeline steps to be processed. imports : set[str] Set of import statements extracted from solutions. instantiations : dict[str, str] Dictionary mapping variable names to instantiation code. pipeline_definition : str The complete generated pipeline code. Examples -------- >>> from bukka.expert_system.solution import Solution >>> from bukka.expert_system.problems import Problem >>> from pathlib import Path >>> >>> # Create a transformer solution >>> scaler = Solution( ... name="scaler", ... function_import="from sklearn.preprocessing import StandardScaler", ... function_name="StandardScaler", ... function_kwargs={} ... ) >>> >>> # Create a problem for the transformer >>> scaling_problem = Problem( ... problem_name="Scaling", ... description="Features need scaling", ... features=["age", "income"], ... solutions=[scaler], ... problem_type="transformer" ... ) >>> >>> # Create pipeline steps >>> pipeline_steps = [(scaler, scaling_problem)] >>> >>> # Write the pipeline >>> writer = PipelineWriter( ... pipeline_steps=pipeline_steps, ... output_path=Path("my_pipeline.py") ... ) >>> writer.write_code() # Generates and writes the pipeline code Notes ----- This class avoids strong typing on solution objects because the project's `Solution` wrappers vary; it performs several runtime checks to extract import strings, instantiation text and step names. """
[docs] def __init__(self, pipeline_steps: list[tuple[Any, Any]], output_path: str | Path) -> None: """Create a PipelineWriter. Parameters ---------- pipeline_steps : list[tuple[Any, Any]] List of (solution_object, problem_object) tuples to include in the pipeline. output_path : str | Path The file path where the generated pipeline code will be written. Examples -------- >>> from pathlib import Path >>> writer = PipelineWriter( ... pipeline_steps=[], ... output_path=Path("empty_pipeline.py") ... ) """ self.pipeline_steps = pipeline_steps if pipeline_steps is not None else [] # Extract all components before initializing parent self.imports: set[str] = set() self.instantiations: dict[str, str] = {} self.pipeline_definition: str = "" # Build all components self._fetch_step_definitions() self._fetch_imports() self._parse_pipeline_steps() # Initialize parent class with template super().__init__( template=FULL_TEMPLATE, output_path=output_path, kwargs=self._build_template_kwargs(), expected_args=["imports", "instantiations", "preprocessor", "pipeline"] ) # Store the pipeline definition for backward compatibility self.pipeline_definition = self._fill_template()
def _fetch_imports(self) -> None: """Populate self.imports from the selected pipeline steps. Examples -------- >>> # Assuming writer has pipeline_steps with solutions >>> writer._fetch_imports() >>> "from sklearn.pipeline import Pipeline" in writer.imports True """ self.imports: set[str] = set() # Add hardcoded imports from the template self.imports.add("from sklearn.pipeline import Pipeline") self.imports.add("from sklearn.compose import ColumnTransformer") for sol_obj, _ in self.pipeline_steps: imp = sol_obj.fetch_import() if imp: self.imports.add(imp) def _fetch_step_definitions(self) -> None: """Create instantiation lines for each pipeline step. The produced strings are simple assignment expressions such as ``step_name = SomeTransformer(arg=val)``. The method is defensive and supports both small wrapper objects and tuples/lists. Ensures unique variable names by appending counters when duplicates are found. Examples -------- >>> # Assuming writer has pipeline_steps >>> writer._fetch_step_definitions() >>> "scaler" in writer.instantiations True >>> "StandardScaler()" in writer.instantiations["scaler"] True """ instantiations: dict[str, str] = {} used_names: set[str] = set() for i, (sol_obj, problem) in enumerate(self.pipeline_steps, start=1): # Decide a reasonable variable/name for the step if hasattr(sol_obj, "name"): name = sol_obj.name else: name = f"step_{i}" # Make a safe python identifier for the variable var_name = self.make_python_string_variable_safe(name, lowercase=True) # Ensure uniqueness by appending a counter if needed counter = 1 while var_name in used_names: var_name = self.make_python_string_variable_safe(f"{var_name}_{counter}", lowercase=True) counter += 1 used_names.add(var_name) # Instantiate using helper method if available inst = sol_obj.fetch_instantiation() instantiations[var_name] = inst self.instantiations = instantiations def _parse_pipeline_steps(self) -> None: self.transformers: list[tuple[str, str, list[str]]] = [] # (name, var_name, columns) self.manipulators: list[str] = [] # var_names for multi-column steps self.model_step: str | None = None for (sol_obj, problem), var_name in zip(self.pipeline_steps, self.instantiations.keys()): problem_type = getattr(problem, "problem_type", None) features = getattr(problem, "features", []) if problem_type == "transformer": # Single-column transformer self.transformers.append((var_name, var_name, features)) elif problem_type == "manipulator": # Multi-column manipulator self.manipulators.append(var_name) elif problem_type == "model": # Final model step self.model_step = var_name else: # Default: treat as manipulator if no type specified self.manipulators.append(var_name) def _build_preprocessor(self) -> str: """Build the ColumnTransformer preprocessor code if needed. This method constructs the ColumnTransformer definition based on the pipeline steps that are transformers. It groups transformers by their target columns and chains multiple transformers on the same columns into a Pipeline. Returns ------- str The ColumnTransformer definition code, or an empty string if no transformers are present. Examples -------- >>> # Assuming writer has pipeline_steps with transformers >>> preprocessor_code = writer._build_preprocessor() >>> "ColumnTransformer" in preprocessor_code True """ if not self.transformers: return "" # Group transformers by their target columns from collections import defaultdict grouped_transformers: dict[tuple[str, ...], list[tuple[str, str]]] = defaultdict(list) for name, var_name, columns in self.transformers: # Convert list to tuple for hashable key col_key = tuple(sorted(columns)) if columns else () grouped_transformers[col_key].append((name, var_name)) ct_items = [] for col_key, transformers in grouped_transformers.items(): columns = list(col_key) if len(transformers) == 1: # Single transformer for these columns name, var_name = transformers[0] ct_items.append(f"('{name}', {var_name}, {repr(columns)})") else: # Multiple transformers for same columns - chain them in a Pipeline pipe_steps = [f"('{name}', {var_name})" for name, var_name in transformers] pipe_steps_str = SEPARATOR.join(pipe_steps) # Use first transformer name as the pipeline name pipeline_name = f"{transformers[0][0]}_pipeline" ct_items.append(f"('{pipeline_name}', Pipeline([{pipe_steps_str}]), {repr(columns)})") return SEPARATOR.join(ct_items) def _build_pipeline_str(self) -> str: """Build the final pipeline definition string. This method constructs the final pipeline definition by combining all pipeline steps, including the preprocessor, manipulators, and model step. Returns ------- str The final pipeline definition string. Examples -------- >>> # Assuming writer has pipeline_steps >>> pipeline_str = writer._build_pipeline_str() >>> "pipeline = Pipeline([" in pipeline_str True """ steps = [] # Add preprocessor if exists if self.transformers: steps.append("('preprocessor', preprocessor)") # Add manipulators for var_name in self.manipulators: steps.append(f"('{var_name}', {var_name})") # Add model as final step if self.model_step: steps.append(f"('{self.model_step}', {self.model_step})") return SEPARATOR.join(steps) def _build_imports_str(self) -> str: """Build the imports section of the pipeline file. Returns ------- str The formatted import statements for the pipeline file. Examples -------- >>> # Assuming writer has imports populated >>> imports_str = writer._build_imports_str() >>> "from sklearn.pipeline import Pipeline" in imports_str True """ if self.imports: return "\n".join(sorted(self.imports)) else: return "" def _build_instantiations_str(self) -> str: """Build the instantiations section of the pipeline file. Returns ------- str The formatted instantiation lines for the pipeline file. Examples -------- >>> # Assuming writer has instantiations populated >>> inst_str = writer._build_instantiations_str() >>> "scaler = StandardScaler()" in inst_str True """ inst_lines = [ f"{var_name} = {inst}" for var_name, inst in self.instantiations.items() ] return "\n".join(inst_lines) if inst_lines else "" def _build_template_kwargs(self) -> dict[str, Any]: """Build the kwargs dictionary for template substitution. Constructs all sections of the pipeline file: imports, instantiations, preprocessor (ColumnTransformer), and final pipeline definition. Returns ------- dict[str, Any] Dictionary with keys 'imports', 'instantiations', 'preprocessor', and 'pipeline' containing the formatted code sections. Examples -------- >>> # Assuming writer is properly initialized >>> kwargs = writer._build_template_kwargs() >>> "imports" in kwargs True >>> "pipeline" in kwargs True """ return { "imports": self._build_imports_str(), "instantiations": self._build_instantiations_str(), "transformer_steps": self._build_preprocessor(), "pipeline": self._build_pipeline_str() }