Add two more examples, Add syntactic sugar for defining repositories

.gitignore

@@ -495,3 +495,9 @@ report.xml
 
 # Auto-generated during builds
 /src/cosy_luigi/_version.py
+
+# Examples
+/examples/**/*.txt
+/examples/**/*.json
+/examples/**/*.csv
+/examples/**/*.skops

README.md

@@ -42,6 +42,16 @@ pip install https://github.com/cls-python/cosy-luigi/releases/download/nightly/c
 ## Documentation
 Please head over to the [documentation](https://cls-python.github.io/cosy-luigi/) to [get started](https://cls-python.github.io/cosy-luigi/quick-start/). 
 
+## Contributing
+Please contribute via a fork if not part of the cls-python org, and contribute via a `feature/` or `bugfix/` branch if you are part of the org. 
+
+Before making a PR:
+- For code, please run `hatch fmt` and `hatch run types:check` to make sure you meet code quality standards. 
+- For docs, please run `hatch run docs:check` to make sure that everything is in order. 
+
+These are run as part of the PR, so if these do not pass for you locally, the PR is guaranteed not to be mergeable. 
+
+
 ## License
 
 `cosy-luigi` is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.

benchmarks/test_benchmark_chain_100.py

@@ -1,11 +1,13 @@
+from abc import ABC
+
 import pytest
 from cosy.maestro import Maestro
 from luigi.mock import MockTarget
 
 from cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
 
 
-class ChainLink(CoSyLuigiTask):
+class ChainLink(CoSyLuigiTask, ABC):
     chain_link: CoSyLuigiTaskParameter | None
 
     def output(self):
@@ -35,7 +37,7 @@ def run(self):
 
 @pytest.fixture
 def repo():
-    return CoSyLuigiRepo(StartingLink, RepeatingLink, FinalLink)
+    return CoSyLuigiRepo(ChainLink, FinalLink)
 
 
 def create_infinite_chain(repo):

examples/getting_started/README.md

@@ -0,0 +1,12 @@
+# Getting Started
+
+This section is still under construction. 
+At present, it contains two examples:
+
+- A very basic example, which primarily demonstrates how Luigi works
+- A simple example on how CoSy-Luigi models variance by inheritance. 
+
+Some interesting quirks: 
+
+- When adding an abstract class or a class that directly inherits from ABC to a repository, it is expanded to all of its concrete implementations. 
+- This behaviour can be manually replicated by calling `get_all_variants` on the class. 

examples/getting_started/basic_example.py

@@ -0,0 +1,48 @@
+import textwrap
+
+import luigi
+from cosy.maestro import Maestro
+
+from src.cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+
+
+class TaskA(CoSyLuigiTask):
+    def output(self):
+        return {"a_artifact": luigi.LocalTarget("output/task_a_output.txt")}
+
+    def run(self):
+        with self.output()["a_artifact"].open("w") as f:
+            f.write("Task A completed")
+
+
+class TaskB(CoSyLuigiTask):
+    task_a = CoSyLuigiTaskParameter(TaskA)
+
+    def output(self):
+        return {"b_artifact": luigi.LocalTarget("output/task_b_output.txt")}
+
+    def run(self):
+        with (
+            self.input()["task_a"]["a_artifact"].open() as input_file,
+            self.output()["b_artifact"].open("w") as output_file,
+        ):
+            data = input_file.read()
+            output_file.write("Task B completed with input: " + data)
+
+
+if __name__ == "__main__":
+    repo = CoSyLuigiRepo(
+        TaskA,
+        TaskB,
+    )
+    maestro = Maestro(repo.cls_repo, repo.taxonomy)
+    results = list(maestro.query(TaskB.target()))
+    luigi.build(results, local_scheduler=True, detailed_summary=True)
+    print(
+        textwrap.dedent(
+            f"""
+                ===============================================
+                    There are a total of {len(results)} results
+                ==============================================="""
+        )
+    )

examples/getting_started/variation_example.py

@@ -0,0 +1,59 @@
+import textwrap
+from abc import ABC
+from string import Template
+
+import luigi
+from cosy.maestro import Maestro
+
+from src.cosy_luigi.combinatorics import CoSyLuigiRepo, CoSyLuigiTask, CoSyLuigiTaskParameter
+
+
+class WriteTemplateTask(CoSyLuigiTask):
+    def output(self):
+        return {"template": luigi.LocalTarget("hello_world_template.txt")}
+
+    def run(self):
+        with self.output()["template"].open("w") as result:
+            result.write("Hello World $name")
+
+
+class SubstituteNameTask(CoSyLuigiTask, ABC):
+    template_task = CoSyLuigiTaskParameter(WriteTemplateTask)
+    name: str = None
+
+    def output(self):
+        return {"filled_template": luigi.LocalTarget(self.__class__.__name__ + "_filled_template.txt")}
+
+    def run(self):
+        with self.input()["template_task"]["template"].open() as input_template:
+            template = Template(input_template.read())
+            result = template.substitute(name=self.name)
+            with self.output()["filled_template"].open("w") as outfile:
+                outfile.write(result)
+
+
+class SubstituteNameByJohnDoeTask(SubstituteNameTask):
+    name = "John Doe"
+
+
+class SubstituteNameByJaneDoeTask(SubstituteNameTask):
+    name = "Jane Doe"
+
+
+def main():
+    repo = CoSyLuigiRepo(WriteTemplateTask, SubstituteNameTask)
+    maestro = Maestro(repo.cls_repo, repo.taxonomy)
+    results = list(maestro.query(SubstituteNameTask.target()))
+    luigi.build(results, local_scheduler=True, detailed_summary=True)
+    print(
+        textwrap.dedent(
+            f"""
+                ===============================================
+                    There are a total of {len(results)} results
+                ==============================================="""
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/lot_sizing/lot_sizing_pipeline.py

@@ -1,5 +1,6 @@
 import json
 import os
+import textwrap
 from abc import ABC, abstractmethod
 from pathlib import Path
 
@@ -38,7 +39,7 @@ def run(self):
             f.write("1, 5, 7, 8, 9, 10, 14, 16, 19, 21, 19, 23, 24, 26, 26, 26, 28, 26, 28, 30")
 
 
-class PredictDemand(CoSyLuigiTask):
+class PredictDemand(CoSyLuigiTask, ABC):
     get_historic_demand = CoSyLuigiTaskParameter(GetHistoricDemand)
     prediction_horizon = 8
     output_filename: str = ""
@@ -155,15 +156,18 @@ def run_optimizer(self, cost, demand):
     repo = CoSyLuigiRepo(
         GetCosts,
         GetHistoricDemand,
-        PredictDemandByAverage,
-        PredictDemandByLinearRegression,
-        OptimizeLotsByLeastUnitCost,
-        OptimizeLotsByGroff,
-        OptimizeLotsByPartPeriod,
-        OptimizeLotsBySilverMeal,
-        OptimizeLotsByWagnerWhitin,
+        PredictDemand,
+        OptimizeLots,
     )
+    print(PredictDemand.get_all_variants())
     maestro = Maestro(repo.cls_repo, repo.taxonomy)
-    for result in maestro.query(OptimizeLots.target()):
-        # print(deps_tree.print_tree(result))
-        luigi.build([result], local_scheduler=True, detailed_summary=True)
+    results = list(maestro.query(OptimizeLots.target()))
+    luigi.build(results, local_scheduler=True, detailed_summary=True)
+    print(
+        textwrap.dedent(
+            f"""
+            ===============================================
+                There are a total of {len(results)} results
+            ==============================================="""
+        )
+    )

examples/ml_blood_sugar_level/ml_blood_sugar_level.py

@@ -1,4 +1,6 @@
 import json
+import textwrap
+from abc import ABC
 from collections.abc import Callable, Iterable, Mapping, Sequence
 from pathlib import Path
 
@@ -28,7 +30,7 @@ def run(self):
             data=np.c_[diabetes["data"], diabetes["target"]], columns=diabetes["feature_names"] + ["target"]
         )
 
-        df.to_json(self.output().path)
+        df.to_json(self.output()["diabetes_data"].path)
 
 
 class TrainTestSplit(CoSyLuigiTask):
@@ -43,18 +45,18 @@ def output(self):
         }
 
     def run(self):
-        data = pd.read_json(self.input()[0].path)
+        data = pd.read_json(self.input()["diabetes"]["diabetes_data"].path)
         x = data.drop(["target"], axis="columns")
         y = data[["target"]]
         x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.33, random_state=42)
 
-        x_train.to_json(self.output()[0].path)
-        x_test.to_json(self.output()[1].path)
-        y_train.to_json(self.output()[2].path)
-        y_test.to_json(self.output()[3].path)
+        x_train.to_json(self.output()["x_train"].path)
+        x_test.to_json(self.output()["x_test"].path)
+        y_train.to_json(self.output()["y_train"].path)
+        y_test.to_json(self.output()["y_test"].path)
 
 
-class FitTransformScaler(CoSyLuigiTask):
+class FitTransformScaler(CoSyLuigiTask, ABC):
     splitted_data = CoSyLuigiTaskParameter(TrainTestSplit)
 
 
@@ -63,7 +65,7 @@ def output(self):
         return {
             "scaled_x_train": luigi.LocalTarget("minmax_scaled_x_train.json"),
             "scaled_x_test": luigi.LocalTarget("minmax_scaled_x_test.json"),
-            "scaler": luigi.LocalTarget("minmax_scaler.json"),
+            "scaler": luigi.LocalTarget("minmax_scaler.skops"),
         }
 
     def run(self):
@@ -80,7 +82,7 @@ def run(self):
         scaled_x_test.to_json(self.output()["scaled_x_test"].path)
 
         with open(self.output()["scaler"].path, "wb") as outfile:
-            json.dump(scaler, outfile)
+            sio.dump(scaler, outfile)
 
 
 class FitTransformRobustScaler(FitTransformScaler):
@@ -107,7 +109,7 @@ def run(self):
             json.dump(scaler, outfile)
 
 
-class TrainRegressionModel(CoSyLuigiTask):
+class TrainRegressionModel(CoSyLuigiTask, ABC):
     scaled_feats = CoSyLuigiTaskParameter(FitTransformScaler)
     splitted_data = CoSyLuigiTaskParameter(TrainTestSplit)
 
@@ -205,9 +207,16 @@ def main():
         EvaluateRegressionModel,
     )
     maestro = Maestro(repo.cls_repo, repo.taxonomy)
-    for result in maestro.query(EvaluateRegressionModel.target()):
-        # print(deps_tree.print_tree(result))
-        luigi.build([result], local_scheduler=True, detailed_summary=True)
+    results = list(maestro.query(EvaluateRegressionModel.target()))
+    luigi.build(results, local_scheduler=True, detailed_summary=True)
+    print(
+        textwrap.dedent(
+            f"""
+                ===============================================
+                    There are a total of {len(results)} results
+                ==============================================="""
+        )
+    )
 
 
 if __name__ == "__main__":

src/cosy_luigi/combinatorics.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
+import inspect
+from abc import ABC
 from collections import defaultdict
 from functools import cache
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import luigi
 from cosy.core import Constructor, SpecificationBuilder
 from luigi.task_register import Register
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Mapping, Sequence
+    from collections.abc import Callable, Iterable, Mapping, Sequence
 
     from cosy.core.synthesizer import Specification
 
@@ -25,6 +27,11 @@ def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         cls.self = cls
 
+    @classmethod
+    @cache
+    def get_all_variants(cls):
+        return set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in c.get_all_variants()])
+
     @classmethod
     @cache
     def get_all_class_attributes(cls):
@@ -86,13 +93,31 @@ def combinator(cls):
 
 
 class CoSyLuigiRepo:
-    def __init__(self, *tasks: type[CoSyLuigiTask]):
+    def __init__(self, *tasks: type[CoSyLuigiTask] | Iterable[type[CoSyLuigiTask]]):
         Register.disable_instance_cache()
-        self.luigi_repo: list[type[CoSyLuigiTask]] = [*tasks]
+
+        # Accepts completely heterogeneous nested collections
+        def flatten(*heterogeneous_task_collection: type[CoSyLuigiTask] | Iterable[type[CoSyLuigiTask]]):
+            return (
+                task
+                for task_or_task_collection in heterogeneous_task_collection
+                for task in (
+                    flatten(*cast("Iterable[type[CoSyLuigiTask]]", task_or_task_collection))
+                    if isinstance(task_or_task_collection, (tuple, list))
+                    else cast("type[CoSyLuigiTask]", task_or_task_collection).get_all_variants()
+                    if inspect.isabstract(task_or_task_collection)
+                    or ABC in cast("type[CoSyLuigiTask]", task_or_task_collection).__bases__
+                    else (task_or_task_collection,)
+                )
+            )
+
+        # This doesn't technically need to unpack as flatten could be typed to accept packed tuples
+        # But performance is equivalent/faster because the first layer doesn't need to be checked this way
+        self.luigi_repo: set[type[CoSyLuigiTask]] = set(flatten(*tasks))
         self.taxonomy: Mapping[str, set[str]] = defaultdict(set)
-        self.cls_repo: list[tuple[str, Callable, Specification]] = []
+        self.cls_repo: set[tuple[str, Callable, Specification]] = set()
         for task in self.luigi_repo:
-            self.cls_repo.append(task.combinator())
+            self.cls_repo.add(task.combinator())
             for tpe in task.mro()[1:]:
                 if issubclass(tpe, CoSyLuigiTask):
                     # Is a subclass of CosyLuigiTask, but a superclass of task