from bukka.utils.files.file_manager import FileManager
CONFIG_TEMPLATE = """
from pathlib import Path
DATAFRAME_BACKEND = '{backend_name}'
# Current Directory
CURRENT_DIR = Path(__file__).parent
# Dataset Paths
TRAIN_DATASET_PATH = CURRENT_DIR / '{train_relative_path}'
TEST_DATASET_PATH = CURRENT_DIR / '{test_relative_path}'
"""
[docs]
class ConfigWriter:
"""
Generates and writes a configuration file for a Bukka project.
This class constructs a Python source file containing configuration settings
for the Bukka project, such as the DataFrame backend to use and optional
MLflow configuration.
Parameters
----------
output_path : str
The file path where the configuration file will be written.
backend_name : str
The name of the DataFrame backend to use with Narwhals (e.g., 'pyarrow', 'modin', 'cudf', 'dask').
file_manager : FileManager
Manager for project file paths and directory structure.
enable_mlflow : bool, optional
Whether to include MLflow configuration (default: False).
mlflow_tracking_uri : str | None, optional
MLflow tracking URI. If None and enable_mlflow is True,
defaults to mlruns directory (default: None).
Examples
--------
>>> writer = ConfigWriter(output_path="config.py", backend_name="pyarrow", file_manager=fm)
>>> writer.write_config() # Writes the config file with pyarrow backend
"""
def __init__(
self,
output_path: str,
backend_name: str,
file_manager: FileManager,
enable_mlflow: bool = False,
mlflow_tracking_uri: str | None = None
):
self.output_path = output_path
self.backend_name = backend_name
self.file_manager = file_manager
self.enable_mlflow = enable_mlflow
self.mlflow_tracking_uri = mlflow_tracking_uri
[docs]
def write_config(self) -> None:
"""
Write the configuration file to the configured output path.
Generates Python source code from the template and writes it to
the specified file. Conditionally includes MLflow configuration.
Examples
--------
>>> writer = ConfigWriter(output_path="config.py", backend_name="pyarrow", file_manager=fm)
>>> writer.write_config() # Writes the config file with pyarrow backend
"""
config_code = CONFIG_TEMPLATE.format(
backend_name=self.backend_name,
train_relative_path=self.file_manager.train_file_relative_path.as_posix(),
test_relative_path=self.file_manager.test_file_relative_path.as_posix()
)
# Append MLflow configuration if enabled
if self.enable_mlflow:
if self.mlflow_tracking_uri is None:
tracking_uri = f"file:///{self.file_manager.mlruns_path.as_posix()}"
else:
tracking_uri = self.mlflow_tracking_uri
mlflow_config = f'''
# MLflow Configuration
MLFLOW_TRACKING_URI = "{tracking_uri}"
MLFLOW_EXPERIMENT_NAME = "{self.file_manager.project_path.name}_experiment"
'''
config_code += mlflow_config
with open(self.output_path, 'w') as file:
file.write(config_code)