Remove ABC's

Author: cpyle0819

aws_doc_sdk_examples_tools/lliam/adapters/repository.py

@@ -1,10 +1,10 @@
 import abc
 from itertools import islice
 from pathlib import Path
-from typing import Any, Generator, Iterable, List, Tuple
+from typing import Any, Dict, Generator, Iterable, List, Tuple
 
 from aws_doc_sdk_examples_tools.doc_gen import DocGen, Example
-from aws_doc_sdk_examples_tools.fs import Fs, PathFs 
+from aws_doc_sdk_examples_tools.fs import Fs, PathFs
 from aws_doc_sdk_examples_tools.lliam.domain.model import Prompt
 from aws_doc_sdk_examples_tools.lliam.shared_constants import BATCH_PREFIX
 
@@ -23,113 +23,68 @@ def batched(iterable: Iterable, n: int) -> Generator[Tuple, Any, None]:
         yield batch
 
 
-class AbstractPromptRepository(abc.ABC):
-    def __init__(self):
-        self.to_write = {}
-
-    def add(self, prompt: Prompt):
-        self._add(prompt)
-
-    def all_all(self, prompts: List[Prompt]):
-        for prompt in prompts:
-            self._add(prompt)
-
-    def batch(self, prompts: List[Prompt]):
-        self._batch(prompts)
-
-    def commit(self): 
-        self._commit()
-
-    def get(self, id: str) -> Prompt:
-        prompt = self._get(id)
-        return prompt
-
-    def get_all(self, ids: List[str]) -> List[Prompt]:
-        prompts = []
-        for id in ids:
-            prompt = self._get(id)
-            prompts.append(prompt)
-        return prompts
-
-    def set_partition(self, name: str):
-        self.partition_name = name
-
-    @property
-    def partition(self):
-        return self.partition_name or ""
+class FsPromptRepository:
+    to_write: Dict[str, str] = {}
 
-    @abc.abstractmethod
-    def _add(self, product: Prompt):
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def _batch(self, prompts: List[Prompt]):
-        raise NotImplementedError
-    
-    @abc.abstractmethod
-    def _commit(self):
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def _get(self, id: str) -> Prompt:
-        raise NotImplementedError
-
-
-class FsPromptRepository(AbstractPromptRepository):
     def __init__(self, fs: Fs = PathFs()):
-        super().__init__()
         self.fs = fs
 
     def rollback(self):
         # TODO: This is not what rollback is for. We should be rolling back any
         # file changes
         self.to_write = {}
 
-    def _add(self, prompt: Prompt):
+    def add(self, prompt: Prompt):
         self.to_write[prompt.id] = prompt.content
 
-    def _batch(self, prompts: List[Prompt]):
+    def all_all(self, prompts: List[Prompt]):
+        for prompt in prompts:
+            self.add(prompt)
+
+    def batch(self, prompts: List[Prompt]):
         for batch_num, batch in enumerate(batched(prompts, DEFAULT_BATCH_SIZE)):
             batch_name = f"{BATCH_PREFIX}{(batch_num + 1):03}"
             for prompt in batch:
                 prompt.id = f"{batch_name}/{prompt.id}"
-                self._add(prompt)
+                self.add(prompt)
 
-    def _commit(self):
-        base_path = (
-            Path(self.partition) if self.partition else Path(".")
-        )
+    def commit(self):
+        base_path = Path(self.partition) if self.partition else Path(".")
 
         for file_path, content in self.to_write.items():
             if content:
                 full_path = base_path / file_path
                 self.fs.mkdir(full_path.parent)
                 self.fs.write(full_path, content)
 
-    def _get(self, id: str):
+    def get(self, id: str):
         return Prompt(id, self.fs.read(Path(id)))
 
+    def get_all(self, ids: List[str]) -> List[Prompt]:
+        prompts = []
+        for id in ids:
+            prompt = self.get(id)
+            prompts.append(prompt)
+        return prompts
 
-class AbstractDocGenRepository(abc.ABC):
-    def get_new_prompts(self, doc_gen_root: str) -> List[Prompt]:
-        return self._get_new_prompts(doc_gen_root)
+    def set_partition(self, name: str):
+        self.partition_name = name
 
-    @abc.abstractmethod
-    def _get_new_prompts(self, doc_gen_root: str) -> List[Prompt]:
-        raise NotImplementedError
+    @property
+    def partition(self):
+        return self.partition_name or ""
 
 
-class FsDocGenRepository(AbstractDocGenRepository):
+class FsDocGenRepository:
     def __init__(self, fs: Fs = PathFs()):
-        super().__init__()
         self.fs = fs
-    
+
     def rollback(self):
         # TODO: This is not what rollback is for. We should be rolling back any
         # file changes
         self._doc_gen = None
 
-    def _get_new_prompts(self, doc_gen_root: str) -> List[Prompt]:
+    def get_new_prompts(self, doc_gen_root: str) -> List[Prompt]:
         # Right now this is the only instance of DocGen used in this Repository,
         # but if that changes we need to move it up.
         self._doc_gen = DocGen.from_root(Path(doc_gen_root), fs=self.fs)

aws_doc_sdk_examples_tools/lliam/domain/commands.py

@@ -9,7 +9,7 @@ class Command:
 
 @dataclass
 class CreatePrompts(Command):
-    doc_gen_root: str 
+    doc_gen_root: str
     system_prompts: List[str]
     out_dir: str
 

aws_doc_sdk_examples_tools/lliam/service_layer/create_prompts.py

@@ -2,14 +2,12 @@
 
 from aws_doc_sdk_examples_tools.lliam.domain.operations import build_ailly_config
 from aws_doc_sdk_examples_tools.lliam.domain.commands import CreatePrompts
-from aws_doc_sdk_examples_tools.lliam.service_layer.unit_of_work import (
-    AbstractUnitOfWork,
-)
+from aws_doc_sdk_examples_tools.lliam.service_layer.unit_of_work import FsUnitOfWork
 
 logger = logging.getLogger(__name__)
 
 
-def create_prompts(cmd: CreatePrompts, uow: AbstractUnitOfWork):
+def create_prompts(cmd: CreatePrompts, uow: FsUnitOfWork):
     with uow:
         system_prompts = uow.prompts.get_all(cmd.system_prompts)
         ailly_config = build_ailly_config(system_prompts)

aws_doc_sdk_examples_tools/lliam/service_layer/unit_of_work.py

@@ -1,52 +1,28 @@
-import abc
-
+from aws_doc_sdk_examples_tools.fs import Fs
 from aws_doc_sdk_examples_tools.lliam.adapters.repository import (
-    AbstractPromptRepository,
-    AbstractDocGenRepository,
     FsPromptRepository,
     FsDocGenRepository,
 )
 
 
-class AbstractUnitOfWork(abc.ABC):
-    prompts: AbstractPromptRepository
-    doc_gen: AbstractDocGenRepository
+class FsUnitOfWork:
+    def __init__(self, fs: Fs):
+        self.fs = fs
 
-    def __enter__(self) -> "AbstractUnitOfWork":
-        return self
+    def __enter__(self):
+        self.prompts = FsPromptRepository(fs=self.fs)
+        self.doc_gen = FsDocGenRepository(fs=self.fs)
 
     def __exit__(self, *args):
         self.rollback()
 
     def commit(self):
-        self._commit()
+        self.prompts.commit()
 
     def collect_new_prompts(self):
         for prompt in self.prompts.seen:
             yield prompt
 
-    @abc.abstractmethod
-    def _commit(self):
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def rollback(self):
-        raise NotImplementedError
-
-
-class FsUnitOfWork(AbstractUnitOfWork):
-
-    def __enter__(self):
-        self.prompts = FsPromptRepository()
-        self.doc_gen = FsDocGenRepository()
-        return super().__enter__()
-
-    def __exit__(self, *args):
-        super().__exit__(*args)
-
-    def _commit(self):
-        self.prompts.commit()
-
     def rollback(self):
         self.prompts.rollback()
         self.doc_gen.rollback()

aws_doc_sdk_examples_tools/lliam/test/services_test.py

@@ -3,41 +3,23 @@
 from aws_doc_sdk_examples_tools.fs import Fs, RecordFs
 from aws_doc_sdk_examples_tools.lliam.domain.commands import CreatePrompts
 from aws_doc_sdk_examples_tools.lliam.service_layer.create_prompts import create_prompts
-from aws_doc_sdk_examples_tools.lliam.service_layer.unit_of_work import AbstractUnitOfWork
-from aws_doc_sdk_examples_tools.lliam.adapters.repository import (
-    FsPromptRepository,
-    FsDocGenRepository,
-)
-
-
-class FakeUnitOfWork(AbstractUnitOfWork):
-    def __init__(self, fs: Fs):
-        self.prompts = FsPromptRepository(fs=fs)
-        self.doc_gen = FsDocGenRepository(fs=fs)
-        self.committed = False
-
-    def _commit(self):
-        self.committed = True
-        self.prompts.commit()
-
-    def rollback(self):
-        self.committed = False
-        self.prompts.rollback()
-        self.doc_gen.rollback()
+from aws_doc_sdk_examples_tools.lliam.service_layer.unit_of_work import FsUnitOfWork
 
 
 def test_create_prompts_writes_when_commit_called():
     """Test that create_prompts successfully writes prompts when commit is called."""
-    fs = RecordFs({
-        Path("system1.md"): "System prompt 1 content",
-        Path("system2.md"): "System prompt 2 content",
-        Path("fake/doc_gen_root"): ""
-    })
-    uow = FakeUnitOfWork(fs=fs)
+    fs = RecordFs(
+        {
+            Path("system1.md"): "System prompt 1 content",
+            Path("system2.md"): "System prompt 2 content",
+            Path("fake/doc_gen_root"): "",
+        }
+    )
+    uow = FsUnitOfWork(fs=fs)
     cmd = CreatePrompts(
         doc_gen_root="fake/doc_gen_root",
         system_prompts=["system1.md", "system2.md"],
-        out_dir="/fake/output"
+        out_dir="/fake/output",
     )
 
     create_prompts(cmd, uow)
@@ -47,4 +29,4 @@ def test_create_prompts_writes_when_commit_called():
     assert fs.stat(ailly_config_path).exists
     ailly_config_content = fs.read(ailly_config_path)
     assert "System prompt 1 content" in ailly_config_content
-    assert "System prompt 2 content" in ailly_config_content
+    assert "System prompt 2 content" in ailly_config_content