51 lines
2.1 KiB
Python
51 lines
2.1 KiB
Python
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from modules.shared.ConfigDataClass import Config
|
|
|
|
|
|
@dataclass
|
|
class ModelPackageCollection(Config):
|
|
name: str = None
|
|
external_packages: list[str] = None
|
|
unsorted_packages: list[str] = None
|
|
categorized_packages: dict[str, list[str]] = None
|
|
|
|
|
|
def __post_init__(self):
|
|
if not self.external_packages: self.external_packages = list()
|
|
if not self.unsorted_packages: self.unsorted_packages = list()
|
|
if not self.categorized_packages: self.categorized_packages = dict()
|
|
super().__post_init__()
|
|
if not self.name: raise ValueError('ModelPackageCollection(): name must be specified')
|
|
|
|
def add_package(self, pkg_name, category: str = None, internal=True):
|
|
if not internal and pkg_name not in self.external_packages: self.external_packages.append(pkg_name)
|
|
elif not category and pkg_name not in self.unsorted_packages: self.unsorted_packages.append(pkg_name)
|
|
else:
|
|
if category not in self.categorized_packages: self.categorized_packages[category] = list()
|
|
if pkg_name not in self.categorized_packages[category]: self.categorized_packages[category].append(pkg_name)
|
|
self.save()
|
|
|
|
def get_path(self, pkg_name) -> Path:
|
|
if pkg_name in self.external_packages: return Path('')
|
|
elif pkg_name in self.unsorted_packages: return Path(self.name)
|
|
else:
|
|
for category in self.categorized_packages:
|
|
if pkg_name in self.categorized_packages[category]: return Path(self.name) / category
|
|
|
|
raise FileNotFoundError(f'package {pkg_name} not in collection {self.name}')
|
|
|
|
@property
|
|
def paths_dict(self) -> dict[str, Path]:
|
|
result: dict[str, Path] = dict()
|
|
for category in self.categorized_packages:
|
|
for pkg_name in list(self.categorized_packages[category]):
|
|
result[pkg_name] = Path(self.name) / category
|
|
|
|
for pkg_name in list(self.unsorted_packages): result[pkg_name] = Path(self.name)
|
|
for pkg_name in list(self.external_packages): result[pkg_name] = Path('')
|
|
|
|
return result
|
|
|