1  from copy import deepcopy
 2  from pathlib import Path
 3  from typing import List, Optional, Self
 4  
 5  from pydantic import BaseModel, Field, PositiveInt, model_validator
 6  
 7  from src.analysis.metrics.config import MetricToScrape, ScrapeConfig
 8  
 9  
10  class PlotConfig(BaseModel):
11      name: Optional[str] = None
12      ignore_columns: List[str] = Field(default_factory=lambda: ["bootstrap", "midstrap"])
13      data_points: PositiveInt = Field(default=25)
14      xlabel_name: str = "Simulation"
15      ylabel_name: str = "KBytes/s"
16      show_min_max: bool = False
17      outliers: bool = True
18      scale_x: PositiveInt = 1000
19      fig_size: List[PositiveInt] = Field(default_factory=lambda: [20, 20])
20  
21      # TODO [plotter config]: Change to paths. Edit plotter too.
22      folder: List[str] = Field(default_factory=list)
23      data: List[str] = Field(default_factory=list)
24      include_files: List[str] = Field(default_factory=list)
25  
26  
27  class PlotConfigBuilder(BaseModel):
28      name: str
29      config: PlotConfig = Field(default=None)
30  
31      @model_validator(mode="after")
32      def _sync_config(self) -> "PlotConfigBuilder":
33          self.config = PlotConfig(name=self.name)
34          return self
35  
36      def with_metric(self, metric: MetricToScrape | str) -> Self:
37          if isinstance(metric, MetricToScrape):
38              self.config.data.append(metric.folder_name.strip("/"))
39          else:
40              self.config.data.append(metric.strip("/"))
41          return self
42  
43      def with_include_files(self, includes: List[str] | str) -> Self:
44          self.config.include_files.extend(includes)
45          return self
46  
47      def with_folders(self, folders: List[str | Path] | str | Path) -> Self:
48          if isinstance(folders, str) or isinstance(folders, Path):
49              folders = [folders]
50  
51          # TODO [plotter config]: This hack will be removed.
52          def ensure_trailing_slash(folder: str | Path) -> str:
53              if isinstance(folder, Path):
54                  folder = folder.as_posix()
55              if folder.endswith("/"):
56                  return folder
57              else:
58                  return f"{folder}/"
59  
60          folders = [ensure_trailing_slash(folder) for folder in folders]
61          self.config.folder.extend(folders)
62          return self
63  
64      def with_scrape_metrics(self, scrape_config: ScrapeConfig) -> Self:
65          for metric in scrape_config.metrics_to_scrape:
66              self.with_metric(metric)
67          return self
68  
69      def with_data_from_scrapes(self, scrape_configs: List[ScrapeConfig] | ScrapeConfig) -> Self:
70          if isinstance(scrape_configs, ScrapeConfig):
71              scrape_configs = [scrape_configs]
72          self.config.include_files.extend([config.name for config in scrape_configs])
73          self.with_folders([config.dump_location for config in scrape_configs])
74          return self
75  
76      def build(self) -> PlotConfig:
77          return deepcopy(self.config)