Simplified code

1 files changed, 14 insertions, 18 deletions

diff --git a/day7/__init__.py b/day7/__init__.py
index 279201a..4adb0ca 100644
--- a/day7/__init__.py
+++ b/day7/__init__.py
@@ -43,7 +43,7 @@ class Folder(Inode):
     def size(self, value: int):
         pass
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Inode]:
         yield self
 
         for inode in self.children.values():
@@ -60,15 +60,14 @@ class File(Inode):
 
 
 @dataclass
-class FileSystem:
-    current: Folder = None
-    root: Folder = field(default_factory=lambda: Folder(name="/"))
+class FileSystem(Folder):
+    pwd: Folder = None
 
     capacity = 70000000
 
     @classmethod
     def from_commands(cls, commands: List[Command]) -> "FileSystem":
-        file_system = cls()
+        file_system = cls(name="/")
 
         for command in commands:
             file_system.parse_command(command)
@@ -76,7 +75,7 @@ class FileSystem:
         return file_system
 
     def __post_init__(self):
-        self.current = self.root
+        self.pwd = self
 
     def parse_command(self, command: Command):
         match command.action:
@@ -90,36 +89,33 @@ class FileSystem:
             size, name = item.split(" ")
 
             inode = (
-                Folder(name=name, parent=self.current)
+                Folder(name=name, parent=self.pwd)
                 if size == "dir"
-                else File(name=name, parent=self.current, size=int(size))
+                else File(name=name, parent=self.pwd, size=int(size))
             )
 
-            if name in self.current.children:
+            if name in self.pwd.children:
                 raise Exception(f"File/Folder already exists {name}")
 
-            self.current.children[name] = inode
+            self.pwd.children[name] = inode
 
     def change_directory(self, name: str):
         match name:
             case "..":
-                new_dir = self.current.parent
+                new_dir = self.pwd.parent
             case "/":
-                new_dir = self.root
+                new_dir = self
             case name:
-                new_dir = self.current.children.get(name)
+                new_dir = self.pwd.children.get(name)
 
         if new_dir is None or isinstance(new_dir, File):
             raise Exception(f"No such folder {name}")
 
-        self.current = new_dir
-
-    def __iter__(self):
-        return iter(self.root)
+        self.pwd = new_dir
 
     @property
     def free_space(self):
-        return self.capacity - self.root.size
+        return self.capacity - self.size
 
 
 class Assignment(BaseAssignment, ABC):