Use case: Ensure files are not imported directly outside specified folder

src/codegen/sdk/core/directory.py

@@ -21,7 +21,7 @@
 
 
 @apidoc
 class Directory(
     HasSymbols[TFile, TSymbol, TImportStatement, TGlobalVar, TClass, TFunction, TImport],
     Generic[TFile, TSymbol, TImportStatement, TGlobalVar, TClass, TFunction, TImport],
 ):
@@ -50,11 +50,28 @@
     def __iter__(self):
         return iter(self.items.values())
 
+    def _is_a_subdirectory_of(self, target_directory: "Directory"):
+        """Checks whether this directory is a subdirectory of another directory"""
+        if self.parent == target_directory:
+            return True
+        if self.parent is None:
+            return False
+        return self.parent._is_a_subdirectory_of(target_directory=target_directory)
+
     def __contains__(self, item: str | TFile | Self) -> bool:
         if isinstance(item, str):
             return item in self.items
+        elif isinstance(item, Directory):
+            return item._is_a_subdirectory_of(self)
         else:
-            return item in self.items.values()
+            # It could only ever be a file here, at least according to item's types...
+            match item.directory:
+                case None:
+                    return False
+                case _ if item.directory == self:
+                    return True
+                case _:
+                    return item.directory._is_a_subdirectory_of(self)
 
     def __len__(self) -> int:
         return len(self.items)
@@ -111,7 +128,7 @@
                     _get_subdirectories(item)
 
         _get_subdirectories(self)
         return subdirectories
 
     @noapidoc
     @cached_generator()
@@ -141,10 +158,10 @@
 
         if ignore_case:
             return next(
                 (f for name, f in self.items.items() if name.lower() == filename.lower() and isinstance(f, File)),
                 None,
             )
         return self.items.get(filename, None)
 
     def add_subdirectory(self, subdirectory: Self) -> None:
         """Add a subdirectory to the directory."""
@@ -163,7 +180,7 @@
 
     def get_subdirectory(self, subdirectory_name: str) -> Self | None:
         """Get a subdirectory by its name (relative to the directory)."""
         return self.items.get(subdirectory_name, None)
 
     def update_filepath(self, new_filepath: str) -> None:
         """Update the filepath of the directory and its contained files."""

tests/unit/codegen/sdk/core/test_directory.py

@@ -234,3 +234,107 @@
         file = codebase.get_file("test/我很喜欢冰激淋/test-file 12'3_🍦.py")
         assert file is not None
         assert file.content == "print('Hello, world!')"
+
+
+def test_contains_dirs_and_files(tmpdir) -> None:
+    # language=python
+    with get_codebase_session(
+        tmpdir=tmpdir,
+        files={
+            "file0.py": "",
+            "main_dir/file1.py": "",
+            "main_dir/file2.py": "",
+            "main_dir/sub_dir/file3.py": "",
+            "main_dir/sub_dir/sub_sub_dir/file4.py": "",
+            "main_dir/sub_dir/sub_sub_dir/file5.py": "",
+            "main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/file6.py": "",
+            "main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/sub_sub_sub_sub_dir/sub_sub_sub_sub_sub_dir/file7.py": "",
+            "lonely_dir/file_lonely.py": "",
+        },
+    ) as codebase:
+        file0 = codebase.get_file("file0.py")
+        main_dir = codebase.get_directory("main_dir")
+        file1 = codebase.get_file("main_dir/file1.py")
+        file2 = codebase.get_file("main_dir/file2.py")
+        sub_dir = codebase.get_directory("main_dir/sub_dir")
+        file3 = codebase.get_file("main_dir/sub_dir/file3.py")
+        sub_sub_dir = codebase.get_directory("main_dir/sub_dir/sub_sub_dir")
+        file4 = codebase.get_file("main_dir/sub_dir/sub_sub_dir/file4.py")
+        file5 = codebase.get_file("main_dir/sub_dir/sub_sub_dir/file5.py")
+        sub_sub_sub_dir = codebase.get_directory("main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir")
+        file6 = codebase.get_file("main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/file6.py")
+        sub_sub_sub_sub_dir = codebase.get_directory("main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/sub_sub_sub_sub_dir")
+        sub_sub_sub_sub_sub_dir = codebase.get_directory("main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/sub_sub_sub_sub_dir/sub_sub_sub_sub_sub_dir")
+        file7 = codebase.get_file("main_dir/sub_dir/sub_sub_dir/sub_sub_sub_dir/sub_sub_sub_sub_dir/sub_sub_sub_sub_sub_dir/file7.py")
+        directory_stack = main_dir.subdirectories
+        directory_stack.append(main_dir)
+        main_directory_stack_no_root = directory_stack
+        file_stack = [file7]
+        for directory in directory_stack:
+            # ignore self
+            if directory != sub_sub_sub_sub_sub_dir:
+                assert sub_sub_sub_sub_sub_dir in directory
+            else:
+                # A dir is not in itself!
+                assert sub_sub_sub_sub_sub_dir not in directory
+            for file in file_stack:
+                assert file in directory
+
+        directory_stack.remove(sub_sub_sub_sub_sub_dir)
+
+        for directory in directory_stack:
+            if directory != sub_sub_sub_sub_dir:
+                assert sub_sub_sub_sub_dir in directory
+
+            for file in file_stack:
+                assert file in directory
+
+        directory_stack.remove(sub_sub_sub_sub_dir)
+        file_stack.append(file6)
+
+        for directory in directory_stack:
+            if directory != sub_sub_sub_dir:
+                assert sub_sub_sub_dir in directory
+
+            for file in file_stack:
+                assert file in directory
+
+        directory_stack.remove(sub_sub_sub_dir)
+        file_stack.append(file5)
+        file_stack.append(file4)
+
+        for directory in directory_stack:
+            if directory != sub_sub_dir:
+                assert sub_sub_dir in directory
+
+            for file in file_stack:
+                assert file in directory
+
+        directory_stack.remove(sub_sub_dir)
+        file_stack.append(file3)
+
+        for directory in directory_stack:
+            if directory != sub_dir:
+                assert sub_dir in directory
+            for file in file_stack:
+                assert file in directory
+
+        directory_stack.remove(sub_dir)
+        file_stack.append(file2)
+        file_stack.append(file1)
+
+        for directory in directory_stack:
+            if directory != main_dir:
+                assert main_dir in directory
+            for file in file_stack:
+                assert file in directory
+
+        lonely_dir = codebase.get_directory("lonely_dir")
+        lonely_file = codebase.get_file("lonely_dir/file_lonely.py")
+
+        for directory in main_directory_stack_no_root:
+            assert file0 not in directory
+            assert lonely_dir not in directory
+            assert lonely_file not in directory
+
+        assert lonely_file in lonely_dir