indent 0.0.8__py3-none-any.whl

exponent/core/remote_execution/files.py

@@ -0,0 +1,415 @@
+import os
+from asyncio import gather, to_thread
+from typing import Final, cast
+
+from anyio import Path as AsyncPath
+from python_ripgrep import PySortMode, PySortModeKind, files, search
+from rapidfuzz import process
+
+from exponent.core.remote_execution.types import (
+    FileAttachment,
+    FilePath,
+    GetAllTrackedFilesRequest,
+    GetAllTrackedFilesResponse,
+    GetFileAttachmentRequest,
+    GetFileAttachmentResponse,
+    GetFileAttachmentsRequest,
+    GetFileAttachmentsResponse,
+    GetMatchingFilesRequest,
+    GetMatchingFilesResponse,
+    ListFilesRequest,
+    ListFilesResponse,
+    RemoteFile,
+)
+from exponent.core.remote_execution.utils import safe_read_file
+
+MAX_MATCHING_FILES: Final[int] = 10
+FILE_NOT_FOUND: Final[str] = "File {} does not exist"
+MAX_FILES_TO_WALK: Final[int] = 10_000
+
+GLOB_MAX_COUNT: Final[int] = 1000
+
+
+class FileCache:
+    """A cache of the files in a working directory.
+
+    Args:
+        working_directory: The working directory to cache the files from.
+    """
+
+    def __init__(self, working_directory: str) -> None:
+        self.working_directory = working_directory
+        self._cache: list[str] | None = None
+
+    async def get_files(self) -> list[str]:
+        """Get the files in the working directory.
+
+        Returns:
+            A list of file paths in the working directory.
+        """
+        if self._cache is None:
+            self._cache = await file_walk(self.working_directory)
+
+        return self._cache
+
+
+async def list_files(list_files_request: ListFilesRequest) -> ListFilesResponse:
+    """Get a list of files in the specified directory.
+
+    Args:
+        list_files_request: An object containing the directory to list files from.
+
+    Returns:
+        A list of RemoteFile objects representing the files in the directory.
+    """
+
+    filenames = [
+        entry.name async for entry in AsyncPath(list_files_request.directory).iterdir()
+    ]
+
+    return ListFilesResponse(
+        files=[
+            RemoteFile(
+                file_path=filename,
+                working_directory=list_files_request.directory,
+            )
+            for filename in filenames
+        ],
+        correlation_id=list_files_request.correlation_id,
+    )
+
+
+async def get_file_content(
+    absolute_path: FilePath, offset: int | None = None, limit: int | None = None
+) -> tuple[str, bool]:
+    """Get the content of the file at the specified path.
+
+    Args:
+        absolute_path: The absolute path to the file.
+
+    Returns:
+        A tuple containing the content of the file and a boolean indicating if the file exists.
+    """
+    file = AsyncPath(absolute_path)
+    exists = await file.exists()
+
+    if not exists:
+        return FILE_NOT_FOUND.format(absolute_path), False
+
+    if await file.is_dir():
+        return "File is a directory", True
+
+    content = await safe_read_file(file)
+
+    if offset or limit:
+        offset = offset or 0
+        limit = limit or -1
+
+        content_lines = content.splitlines()
+        content_lines = content_lines[offset:]
+        content_lines = content_lines[:limit]
+
+        content = "\n".join(content_lines)
+
+    return content, exists
+
+
+async def get_file_attachments(
+    get_file_attachments_request: GetFileAttachmentsRequest,
+    client_working_directory: str,
+) -> GetFileAttachmentsResponse:
+    """Get the content of the files at the specified paths.
+
+    Args:
+        get_file_attachments_request: An object containing the file paths.
+        client_working_directory:  The working directory of the client.
+
+    Returns:
+        A list of FileAttachment objects containing the content of the files.
+    """
+    remote_files = get_file_attachments_request.files
+    attachments = await gather(
+        *[
+            get_file_content(
+                AsyncPath(client_working_directory) / remote_file.file_path
+            )
+            for remote_file in remote_files
+        ]
+    )
+
+    files = [
+        FileAttachment(attachment_type="file", file=remote_file, content=content)
+        for remote_file, (content, _) in zip(remote_files, attachments)
+    ]
+
+    return GetFileAttachmentsResponse(
+        correlation_id=get_file_attachments_request.correlation_id,
+        file_attachments=files,
+    )
+
+
+async def get_file_attachment(
+    get_file_attachment_request: GetFileAttachmentRequest, client_working_directory: str
+) -> GetFileAttachmentResponse:
+    """Get the content of the file at the specified path.
+
+    Args:
+        get_file_attachment_request: An object containing the file path.
+        client_working_directory: The working directory of the client.
+
+    Returns:
+        A FileAttachment object containing the content of the file.
+    """
+    file = get_file_attachment_request.file
+    absolute_path = await file.resolve(client_working_directory)
+
+    content, exists = await get_file_content(absolute_path)
+
+    return GetFileAttachmentResponse(
+        content=content,
+        exists=exists,
+        file=file,
+        correlation_id=get_file_attachment_request.correlation_id,
+    )
+
+
+async def get_matching_files(
+    search_term: GetMatchingFilesRequest,
+    file_cache: FileCache,
+) -> GetMatchingFilesResponse:
+    """Get the files that match the search term.
+
+    Args:
+        search_term: The search term to match against the files.
+        file_cache: A cache of the files in the working directory.
+
+    Returns:
+        A list of RemoteFile objects that match the search term.
+    """
+    # Use rapidfuzz to find the best matching files
+    matching_files = await to_thread(
+        process.extract,
+        search_term.search_term,
+        await file_cache.get_files(),
+        limit=MAX_MATCHING_FILES,
+        score_cutoff=0,
+    )
+
+    directory = file_cache.working_directory
+    files: list[RemoteFile] = [
+        RemoteFile(file_path=file, working_directory=directory)
+        for file, _, _ in matching_files
+    ]
+
+    return GetMatchingFilesResponse(
+        files=files,
+        correlation_id=search_term.correlation_id,
+    )
+
+
+async def get_all_tracked_files(
+    request: GetAllTrackedFilesRequest,
+    working_directory: str,
+) -> GetAllTrackedFilesResponse:
+    return GetAllTrackedFilesResponse(
+        correlation_id=request.correlation_id,
+        files=await get_all_non_ignored_files(working_directory),
+    )
+
+
+async def search_files(
+    path_str: str,
+    file_pattern: str | None,
+    regex: str,
+    working_directory: str,
+) -> list[str]:
+    path = AsyncPath(working_directory) / path_str
+    path_resolved = await path.resolve()
+    globs = [file_pattern] if file_pattern else None
+
+    return await to_thread(
+        search,
+        patterns=[regex],
+        paths=[str(path_resolved)],
+        globs=globs,
+        after_context=3,
+        before_context=5,
+        heading=True,
+        separator_field_context="|",
+        separator_field_match="|",
+        separator_context="\n...\n",
+    )
+
+
+async def get_all_file_contents(
+    working_directory: str,
+) -> list[list[str]]:
+    path_resolved = await AsyncPath(working_directory).resolve()
+
+    results = await to_thread(
+        search,
+        patterns=[".*"],
+        paths=[str(path_resolved)],
+        globs=["!**/poetry.lock", "!**/pnpm-lock.yaml"],
+        heading=True,
+        line_number=False,
+    )
+
+    result_sizes = [len(result) for result in results]
+    total_size = sum(result_sizes)
+    batch_size = total_size // 10
+
+    batches = []
+    current_batch: list[str] = []
+    current_size = 0
+
+    for i, result in enumerate(results):
+        if current_size + result_sizes[i] > batch_size:
+            batches.append(current_batch)
+            current_batch = []
+            current_size = 0
+
+        current_batch.append(result)
+        current_size += result_sizes[i]
+
+    batches.append(current_batch)
+
+    return batches
+
+
+async def normalize_files(
+    working_directory: str, file_paths: list[FilePath]
+) -> list[RemoteFile]:
+    """Normalize file paths to be relative to the working directory.
+
+    Args:
+        working_directory: The working directory to normalize the file paths against.
+        file_paths: A list of file paths to normalize.
+
+    Returns:
+        A list of RemoteFile objects with normalized file paths.
+    """
+    working_path = await AsyncPath(working_directory).resolve()
+    normalized_files = []
+
+    for file_path in file_paths:
+        path = AsyncPath(file_path)
+
+        if path.is_absolute():
+            path = path.relative_to(working_path)
+
+        normalized_files.append(
+            RemoteFile(
+                file_path=str(path),
+                working_directory=working_directory,
+            )
+        )
+
+    return sorted(normalized_files)
+
+
+def _format_ignore_globs(ignore_extra: list[str] | None) -> list[str]:
+    if ignore_extra is None:
+        return []
+
+    return [f"!**/{ignore}" for ignore in ignore_extra]
+
+
+async def file_walk(
+    directory: str,
+    ignore_extra: list[str] | None = None,
+    max_files: int = MAX_FILES_TO_WALK,
+) -> list[str]:
+    """
+    Walk through a directory and return all file paths, respecting .gitignore and additional ignore patterns.
+
+    Args:
+        directory: The directory to walk through
+        ignore_extra: Additional directory paths to ignore, follows the gitignore format.
+        max_files: The maximal number of files to return
+
+    Returns:
+        A list of file paths in the directory.
+    """
+    working_path = str(await AsyncPath(directory).resolve())
+
+    results: list[str] = await to_thread(
+        files,
+        patterns=[""],
+        paths=[working_path],
+        globs=_format_ignore_globs(ignore_extra),
+        sort=PySortMode(kind=PySortModeKind.Path),
+        max_count=max_files,
+    )
+
+    # Create relative paths using os.path functions which handle platform differences
+    relative_results = []
+    for result in results:
+        # Check if the path is inside the working directory
+        if os.path.commonpath([working_path, result]) == working_path:
+            # Create relative path
+            rel_path = os.path.relpath(result, working_path)
+            relative_results.append(rel_path)
+        else:
+            # Fallback to just using the filename
+            relative_results.append(os.path.basename(result))
+
+    return relative_results
+
+
+async def get_all_non_ignored_files(working_directory: str) -> list[RemoteFile]:
+    file_paths = await file_walk(working_directory, ignore_extra=DEFAULT_IGNORES)
+
+    return await normalize_files(working_directory, cast(list[FilePath], file_paths))
+
+
+async def glob(
+    path: str,
+    glob_pattern: str,
+) -> list[str]:
+    return await to_thread(
+        files,
+        patterns=[],
+        paths=[path],
+        globs=[glob_pattern],
+        sort=PySortMode(kind=PySortModeKind.Path),
+        max_count=GLOB_MAX_COUNT,
+    )
+
+
+DEFAULT_IGNORES = [
+    "**/.git/",
+    ".venv/",
+    ".mypy_cache",
+    ".pytest_cache",
+    "node_modules/",
+    "venv/",
+    ".pyenv",
+    "__pycache__",
+    ".ipynb_checkpoints",
+    ".vercel",
+    "__pycache__/",
+    "*.py[cod]",
+    "*$py.class",
+    ".env",
+    "*.so",
+    ".Python",
+    "build/",
+    "develop-eggs/",
+    "dist/",
+    "downloads/",
+    "eggs/",
+    ".eggs/",
+    "lib/",
+    "lib64/",
+    "parts/",
+    "sdist/",
+    "var/",
+    "wheels/",
+    "pip-wheel-metadata/",
+    "share/python-wheels/",
+    "*.egg-info/",
+    ".installed.cfg",
+    "*.egg",
+    "MANIFEST",
+]

exponent/core/remote_execution/git.py

@@ -0,0 +1,268 @@
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import cast
+
+import pygit2
+from anyio import Path as AsyncPath
+from gitignore_parser import (
+    IgnoreRule,
+    handle_negation,
+    parse_gitignore,
+    rule_from_pattern,
+)
+from pygit2 import Tree
+from pygit2.enums import DiffOption
+from pygit2.repository import Repository
+
+from exponent.core.remote_execution.types import (
+    GitInfo,
+)
+from exponent.core.remote_execution.utils import safe_read_file
+
+GIT_OBJ_COMMIT = 1
+
+
+async def git_file_walk(
+    repo: Repository,
+    directory: str,
+) -> list[str]:
+    """
+    Walk through a directory and return all file paths, respecting .gitignore and additional ignore patterns.
+    """
+    tree = get_git_subtree_for_dir(repo, directory)
+
+    if not tree:
+        return []
+
+    # diff to the empty tree to see all files
+    tracked_diff = tree.diff_to_tree()
+
+    tracked_files = [delta.new_file.path for delta in tracked_diff.deltas]
+
+    # Find untracked files relative to the root
+    untracked_diff = repo.diff(flags=DiffOption.INCLUDE_UNTRACKED)
+    untracked_files_from_root = [
+        AsyncPath(delta.new_file.path) for delta in untracked_diff.deltas
+    ]
+
+    # Current working directory relative to the repo root
+    dir_path = await AsyncPath(directory).resolve()
+    repo_path = await AsyncPath(repo.workdir).resolve()
+
+    if repo_path == dir_path:
+        relative_directory = str(repo_path)
+    else:
+        relative_directory = str(dir_path.relative_to(repo_path))
+
+    # Resolve all untracked files that are within the current working directory
+    untracked_files = []
+    for untracked_file in untracked_files_from_root:
+        if not untracked_file.is_relative_to(relative_directory):
+            continue
+
+        untracked_files.append(str(untracked_file.relative_to(relative_directory)))
+
+    # Combine both as sets to remove duplicates
+    return list(set(tracked_files) | set(untracked_files))
+
+
+def get_repo(working_directory: str) -> Repository | None:
+    try:
+        return Repository(working_directory)
+    except pygit2.GitError:
+        return None
+
+
+async def get_git_info(working_directory: str) -> GitInfo | None:
+    try:
+        repo = Repository(working_directory)
+    except pygit2.GitError:
+        return None
+
+    return GitInfo(
+        branch=(await _get_git_branch(repo)) or "<unknown branch>",
+        remote=_get_git_remote(repo),
+    )
+
+
+def get_tracked_files_in_dir(
+    repo: Repository,
+    dir: str | Path,
+    filter_func: Callable[[str], bool] | None = None,
+) -> list[str]:
+    rel_path = get_path_relative_to_repo_root(repo, dir)
+    dir_tree = get_git_subtree_for_dir(repo, dir)
+    entries: list[str] = []
+    if not dir_tree:
+        return entries
+    for entry in dir_tree:
+        if not entry.name:
+            continue
+        entry_path = str(Path(f"{repo.workdir}/{rel_path}/{entry.name}"))
+        if entry.type_str == "tree":
+            entries.extend(get_tracked_files_in_dir(repo, entry_path, filter_func))
+        elif entry.type_str == "blob":
+            if not filter_func or filter_func(entry.name):
+                entries.append(entry_path)
+    return entries
+
+
+def get_git_subtree_for_dir(repo: Repository, dir: str | Path) -> Tree | None:
+    rel_path = get_path_relative_to_repo_root(repo, dir)
+
+    try:
+        head_commit = repo.head.peel(GIT_OBJ_COMMIT)
+    except pygit2.GitError:
+        # If the repo is empty, then the head commit will not exist
+        return None
+    head_tree: Tree = head_commit.tree
+
+    if rel_path == Path("."):
+        # If the relative path is the root of the repo, then
+        # the head_tree is what we want. Note we do this because
+        # Passing "." or "" as the path into the tree will raise.
+        return head_tree
+    return cast(Tree, head_tree[str(rel_path)])
+
+
+def get_path_relative_to_repo_root(repo: Repository, path: str | Path) -> Path:
+    path = Path(path).resolve()
+    return path.relative_to(Path(repo.workdir).resolve())
+
+
+def get_local_commit_hash() -> str:
+    try:
+        # Open the repository (assumes the current working directory is within the git repo)
+        repo = Repository(os.getcwd())
+
+        # Get the current HEAD commit
+        head = repo.head
+
+        # Get the commit object and return its hash as a string
+        return str(repo[head.target].id)
+    except pygit2.GitError:
+        return "unknown-local-commit"
+
+
+def _get_git_remote(repo: Repository) -> str | None:
+    if repo.remotes:
+        return str(repo.remotes[0].url)
+    return None
+
+
+async def _get_git_branch(repo: Repository) -> str | None:
+    try:
+        # Look for HEAD file in the .git directory
+        head_path = AsyncPath(os.path.join(repo.path, "HEAD"))
+
+        if not await head_path.exists():
+            return None
+
+        head_content_raw = await safe_read_file(head_path)
+        head_content = head_content_raw.strip()
+
+        if head_content.startswith("ref:"):
+            return head_content.split("refs/heads/")[-1]
+        else:
+            return None
+
+    except Exception:  # noqa: BLE001
+        return None
+
+
+class GitIgnoreHandler:
+    def __init__(
+        self, working_directory: str, default_ignores: list[str] | None = None
+    ):
+        self.checkers = {}
+
+        if default_ignores:
+            self.checkers[working_directory] = self._parse_ignore_extra(
+                working_directory, default_ignores
+            )
+
+    async def read_ignorefile(self, path: str) -> None:
+        new_ignore = await self._get_ignored_checker(path)
+
+        if new_ignore:
+            self.checkers[path] = new_ignore
+
+    def filter(
+        self,
+        relpaths: list[str],
+        root: str,
+    ) -> list[str]:
+        result = []
+
+        for relpath in relpaths:
+            if relpath.startswith(".git"):
+                continue
+
+            path = os.path.join(root, relpath)
+
+            if self.is_ignored(path):
+                continue
+
+            result.append(relpath)
+
+        return result
+
+    def is_ignored(self, path: str) -> bool:
+        return any(
+            self.checkers[dp](path)
+            for dp in self.checkers
+            if self._is_subpath(path, dp)
+        )
+
+    def _parse_ignore_extra(
+        self, working_directory: str, ignore_extra: list[str]
+    ) -> Callable[[str], bool]:
+        rules: list[IgnoreRule] = []
+
+        for pattern in ignore_extra:
+            if (
+                rule := rule_from_pattern(pattern, base_path=working_directory)
+            ) is not None:
+                rules.append(rule)
+
+        def rule_handler(file_path: str) -> bool:
+            nonlocal rules
+            return bool(handle_negation(file_path, rules))
+
+        return rule_handler
+
+    async def _get_ignored_checker(self, dir_path: str) -> Callable[[str], bool] | None:
+        new_ignore = self._parse_gitignore(dir_path)
+
+        existing_ignore = self.checkers.get(dir_path)
+
+        if existing_ignore and new_ignore:
+            return self._or(new_ignore, existing_ignore)
+
+        return new_ignore or existing_ignore
+
+    @staticmethod
+    def _parse_gitignore(directory: str) -> Callable[[str], bool] | None:
+        gitignore_path = os.path.join(directory, ".gitignore")
+
+        if os.path.isfile(gitignore_path):
+            return cast(Callable[[str], bool], parse_gitignore(gitignore_path))
+
+        return None
+
+    @staticmethod
+    def _or(
+        a: Callable[[str], bool], b: Callable[[str], bool]
+    ) -> Callable[[str], bool]:
+        def or_handler(file_path: str) -> bool:
+            return a(file_path) or b(file_path)
+
+        return or_handler
+
+    @staticmethod
+    def _is_subpath(path: str, parent: str) -> bool:
+        """
+        Check if a path is a subpath of another path.
+        """
+        return os.path.commonpath([path, parent]) == parent