feat: [CG-10650] codebase.codeowners interface

ruff.toml

@@ -57,6 +57,7 @@ extend-generics = [
   "codegen.sdk.core.assignment.Assignment",
   "codegen.sdk.core.class_definition.Class",
   "codegen.sdk.core.codebase.Codebase",
+  "codegen.sdk.core.codeowner.CodeOwner",
   "codegen.sdk.core.dataclasses.usage.Usage",
   "codegen.sdk.core.dataclasses.usage.UsageType",
   "codegen.sdk.core.dataclasses.usage.UsageKind",

src/codegen/git/utils/cache_utils.py

@@ -0,0 +1,43 @@
+import functools
+from collections.abc import Iterator
+from typing import Callable, Generic, ParamSpec, TypeVar
+
+ItemType = TypeVar("ItemType")
+GenParamSpec = ParamSpec("GenParamSpec")
+
+
+class LazyGeneratorCache(Generic[ItemType]):
+    """A cache for a generator that is lazily evaluated."""
+
+    _cache: list[ItemType]
+    gen: Iterator[ItemType]
+
+    def __init__(self, gen: Iterator[ItemType]):
+        self._cache = []
+        self.gen = gen
+
+    def __iter__(self) -> Iterator[ItemType]:
+        for item in self._cache:
+            yield item
+
+        for item in self.gen:
+            self._cache.append(item)
+            yield item
+
+
+def cached_generator(maxsize: int = 16, typed: bool = False) -> Callable[[Callable[GenParamSpec, Iterator[ItemType]]], Callable[GenParamSpec, Iterator[ItemType]]]:
+    """Decorator to cache the output of a generator function.
+
+    The generator's output is fully consumed on the first call and stored as a list.
+    Subsequent calls with the same arguments yield values from the cached list.
+    """
+
+    def decorator(func: Callable[GenParamSpec, Iterator[ItemType]]) -> Callable[GenParamSpec, Iterator[ItemType]]:
+        @functools.lru_cache(maxsize=maxsize, typed=typed)
+        @functools.wraps(func)
+        def wrapper(*args: GenParamSpec.args, **kwargs: GenParamSpec.kwargs) -> Iterator[ItemType]:
+            return LazyGeneratorCache(func(*args, **kwargs))
+
+        return wrapper
+
+    return decorator

src/codegen/sdk/core/codebase.py

@@ -7,6 +7,7 @@
 import re
 from collections.abc import Generator
 from contextlib import contextmanager
+from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Generic, Literal, TypeVar, Unpack, overload
 
@@ -35,6 +36,7 @@
 from codegen.sdk.codebase.span import Span
 from codegen.sdk.core.assignment import Assignment
 from codegen.sdk.core.class_definition import Class
+from codegen.sdk.core.codeowner import CodeOwner
 from codegen.sdk.core.detached_symbols.code_block import CodeBlock
 from codegen.sdk.core.detached_symbols.parameter import Parameter
 from codegen.sdk.core.directory import Directory
@@ -256,6 +258,17 @@ def files(self, *, extensions: list[str] | Literal["*"] | None = None) -> list[T
         # Sort files alphabetically
         return sort_editables(files, alphabetical=True, dedupe=False)
 
+    @cached_property
+    def codeowners(self) -> list["CodeOwner[TSourceFile]"]:
+        """List all CodeOnwers in the codebase.
+
+        Returns:
+            list[CodeOwners]: A list of CodeOwners objects in the codebase.
+        """
+        if self.G.codeowners_parser is None:
+            return []
+        return CodeOwner.from_parser(self.G.codeowners_parser, lambda *args, **kwargs: self.files(*args, **kwargs))
+
     @property
     def directories(self) -> list[TDirectory]:
         """List all directories in the codebase.

src/codegen/sdk/core/codeowner.py

@@ -0,0 +1,75 @@
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Callable, Generic, Literal, ParamSpec, TypeVar
+
+from codeowners import CodeOwners as CodeOwnersParser
+
+from codegen.git.utils.cache_utils import cached_generator
+from codegen.shared.decorators.docs import apidoc
+
+if TYPE_CHECKING:
+    from codegen.sdk.core.file import SourceFile
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+TSourceFile = TypeVar("TSourceFile", bound="SourceFile")
+SourceParam = ParamSpec("SourceParam")
+
+
+@apidoc
+class CodeOwner(Generic[TSourceFile]):
+    """CodeOwner is a class that represents a code owner in a codebase.
+
+    It is used to iterate over all files that are owned by a specific owner.
+
+    Attributes:
+        parser: The CodeOwnersParser that was used to parse the codeowners file.
+        owner_type: The type of the owner (USERNAME, TEAM, EMAIL).
+        owner_value: The value of the owner.
+        files_source: A callable that returns an iterable of all files in the codebase.
+    """
+
+    owner_type: Literal["USERNAME", "TEAM", "EMAIL"]
+    owner_value: str
+    files_source: Callable[SourceParam, Iterable[TSourceFile]]
+
+    def __init__(self, files_source: Callable[SourceParam, Iterable[TSourceFile]], owner_type: Literal["USERNAME", "TEAM", "EMAIL"], owner_value: str):
+        self.owner_type = owner_type
+        self.owner_value = owner_value
+        self.files_source = files_source
+
+    @classmethod
+    def from_parser(cls, parser: CodeOwnersParser, file_source: Callable[SourceParam, Iterable[TSourceFile]]) -> list["CodeOwner"]:
+        """Create a list of CodeOwner objects from a CodeOwnersParser.
+
+        Args:
+            parser (CodeOwnersParser): The CodeOwnersParser to use.
+            file_source (Callable[SourceParam, Iterable[TSourceFile]]): A callable that returns an iterable of all files in the codebase.
+
+        Returns:
+            list[CodeOwner]: A list of CodeOwner objects.
+        """
+        codeowners = []
+        for _, _, owners, _, _ in parser.paths:
+            for owner_label, owner_value in owners:
+                codeowners.append(CodeOwner(file_source, owner_label, owner_value))
+        return codeowners
+
+    @cached_generator(maxsize=16)
+    def files(self, *args: SourceParam.args, **kwargs: SourceParam.kwargs) -> Iterable[TSourceFile]:
+        for source_file in self.files_source(*args, **kwargs):
+            # Filter files by owner value
+            if self.owner_value in source_file.owners:
+                yield source_file
+
+    @property
+    def name(self) -> str:
+        return self.owner_value
+
+    def __iter__(self) -> Iterator[TSourceFile]:
+        return self.files()
+
+    def __repr__(self) -> str:
+        return f"CodeOwner(owner_type={self.owner_type}, owner_value={self.owner_value})"

src/codegen/shared/compilation/function_imports.py

@@ -28,6 +28,7 @@ def get_generated_imports():
 from codegen.sdk.core.codebase import CodebaseType
 from codegen.sdk.core.codebase import PyCodebaseType
 from codegen.sdk.core.codebase import TSCodebaseType
+from codegen.sdk.core.codeowner import CodeOwner
 from codegen.sdk.core.dataclasses.usage import Usage
 from codegen.sdk.core.dataclasses.usage import UsageKind
 from codegen.sdk.core.dataclasses.usage import UsageType