fix(dl): Fixed dl dependencies.

- Added source information for PyTorch dependencies. Now installs CUDA with PyTorch.
- Added ignore logic to stop PyTorch toml info appearing in non-dl projects.

Author: Achronus

zenforge/__init__.py

@@ -20,6 +20,14 @@ class ProjectType(StrEnum):
     API_DEEP_LEARNING = "api-dl"
     ALL = "all"
 
+    @staticmethod
+    def dl_project(value: str) -> bool:
+        """Checks if the value is a Deep Learning project."""
+        if "dl" in value or value == "all":
+            return True
+
+        return False
+
 
 __all__ = [
     "ProjectType",

zenforge/config/toml.py

@@ -1,6 +1,6 @@
-from typing import Any, override
+from typing import Any, Dict, Set, override
 
-TomlSettingsType = dict[str, dict[str, Any]]
+TomlSettingsType = Dict[str, Dict[str, Any]]
 
 
 class TomlOptions:
@@ -12,6 +12,38 @@ def as_dict(self) -> TomlSettingsType:
         return {f"[{self.toml_name}]": self.__dict__}
 
 
+class PoetryDependenciesExtras(TomlOptions):
+    """
+    Stores extra information for poetry dependencies for the `pyproject.toml` file.
+
+    Used for installing cuda with `torch` and `torchvision`.
+    """
+
+    toml_name = "tool.poetry.dependencies"
+
+    def __init__(self) -> None:
+        self.items = {
+            "torch": {"source": "pytorch"},
+            "torchvision": {"source": "pytorch"},
+        }
+
+    def items_to_str_dict(self) -> Dict[str, Set[str]]:
+        """Converts the values of the `self.items` dictionary into a strings."""
+        str_dict = {}
+        for key, d in self.items.items():
+            value_str = []
+            for k, v in d.items():
+                value_str.append(f'{k} = "{v}"')
+
+            str_dict[key] = {", ".join(value_str)}
+
+        return str_dict
+
+    @override
+    def as_dict(self) -> TomlSettingsType:
+        return {f"[{self.toml_name}]": self.items_to_str_dict()}
+
+
 class PyTorchSource(TomlOptions):
     """
     Stores the PyTorch Poetry source details for the `pyproject.toml` file.
@@ -105,7 +137,9 @@ def settings_to_toml(d: TomlSettingsType) -> str:
         toml_str += f"{header}\n"
 
         for key, value in settings.items():
-            if isinstance(value, str):
+            if isinstance(value, set):
+                value = str(value).replace("'", " ")
+            elif isinstance(value, str):
                 value = f'"{value}"'
             elif isinstance(value, bool):
                 value = str(value).lower()
@@ -125,14 +159,17 @@ def multi_settings_to_toml(d_list: list[TomlOptions]) -> str:
     return full
 
 
-def set_toml_settings(project_name: str) -> str:
+def set_toml_settings(project_name: str, dl_project: bool) -> str:
     """Sets the extra toml settings to add to the end of the `pyproject.toml` file."""
     items = [
-        PyTorchSource(),
         PytestOptions(project_name),
         MypyOptions(),
         IsortOptions(),
         BlackOptions(),
     ]
 
+    if dl_project:
+        items.insert(0, PoetryDependenciesExtras())
+        items.insert(1, PyTorchSource())
+
     return multi_settings_to_toml(items)

zenforge/create/method.py

@@ -108,7 +108,10 @@ def __init__(
         self.path = path
         self.deps = deps
 
-        self.toml_extras = set_toml_settings(project_name)
+        self.toml_extras = set_toml_settings(
+            project_name,
+            ProjectType.dl_project(project_type),
+        )
 
         self.env_name = "venv"
         self.python_path = (