Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 17 additions & 5 deletions infrahub_sdk/file_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,11 @@
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING, BinaryIO, cast, overload
from typing import TYPE_CHECKING, BinaryIO, overload

import anyio
import httpx
from anyio.to_thread import run_sync as run_sync_in_thread

from .exceptions import AuthenticationError, NodeNotFoundError, ServerNotReachableError

Expand All @@ -21,6 +22,17 @@ class PreparedFile:
should_close: bool


def _open_binary(path: Path) -> BinaryIO:
"""Open a file in binary read mode.

Wrapper exists to pin the return type to BinaryIO when called via
``run_sync_in_thread``: the overload resolution on ``Path.open``'s mode arg
is lost through the indirection and would otherwise widen to a union of all
open() return types.
"""
return path.open("rb")


class FileHandlerBase:
"""Base class for file handling operations.

Expand Down Expand Up @@ -57,11 +69,11 @@ async def prepare_upload(content: bytes | Path | BinaryIO | None, name: str | No
if isinstance(content, Path):
# Open file in thread pool to avoid blocking the event loop
# Returns a sync file handle that httpx can stream from in chunks
file_obj = await anyio.to_thread.run_sync(content.open, "rb")
return PreparedFile(file_object=cast("BinaryIO", file_obj), filename=filename, should_close=True)
file_obj = await run_sync_in_thread(_open_binary, content)
return PreparedFile(file_object=file_obj, filename=filename, should_close=True)

# At this point, content must be a BinaryIO (file-like object)
return PreparedFile(file_object=cast("BinaryIO", content), filename=filename, should_close=False)
return PreparedFile(file_object=content, filename=filename, should_close=False)

@staticmethod
def prepare_upload_sync(content: bytes | Path | BinaryIO | None, name: str | None = None) -> PreparedFile:
Expand Down Expand Up @@ -92,7 +104,7 @@ def prepare_upload_sync(content: bytes | Path | BinaryIO | None, name: str | Non
return PreparedFile(file_object=content.open("rb"), filename=filename, should_close=True)

# At this point, content must be a BinaryIO (file-like object)
return PreparedFile(file_object=cast("BinaryIO", content), filename=filename, should_close=False)
return PreparedFile(file_object=content, filename=filename, should_close=False)

@staticmethod
def handle_error_response(exc: httpx.HTTPStatusError) -> None:
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -143,10 +143,10 @@ include = ["infrahub_sdk/checks.py"]
invalid-await = "ignore" # 1 violation

[[tool.ty.overrides]]
include = ["infrahub_sdk/file_handler.py", "infrahub_sdk/utils.py"]
include = ["infrahub_sdk/utils.py"]

[tool.ty.overrides.rules]
unresolved-attribute = "ignore" # 5 violations total (1 in file_handler.py, 4 in utils.py)
unresolved-attribute = "ignore" # 4 violations in utils.py

[[tool.ty.overrides]]
include = ["infrahub_sdk/transfer/**"]
Expand Down
51 changes: 50 additions & 1 deletion tests/unit/sdk/test_file_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import tempfile
from io import BytesIO
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, BinaryIO

import anyio
import httpx
Expand All @@ -22,6 +22,15 @@
NODE_ID = "test-node-123"


def _open_binary_handle(path: str) -> BinaryIO:
"""Sync helper to open a file in binary read mode.

Defined as a sync function so ruff's ASYNC* rules don't apply when called from async tests.
The handle is the realistic non-BytesIO BinaryIO that production callers pass to ``prepare_upload``.
"""
return Path(path).open("rb")


async def test_prepare_upload_with_bytes() -> None:
"""Test preparing upload with bytes content (async)."""
content = b"test file content"
Expand Down Expand Up @@ -85,6 +94,26 @@ async def test_prepare_upload_with_binary_io() -> None:
assert prepared.should_close is False


async def test_prepare_upload_with_open_file() -> None:
"""Test preparing upload with a real file handle (not BytesIO) — async.

Locks in BinaryIO branch behaviour for callers passing the result of opening a real
file, not just ``BytesIO``. Guards against narrowing the dispatch to a too-specific
type.
"""
with tempfile.NamedTemporaryFile(suffix=".txt") as tmp:
tmp.write(b"content from open file")
tmp.flush()

with _open_binary_handle(tmp.name) as file_handle:
prepared = await FileHandlerBase.prepare_upload(content=file_handle, name="from_open.bin")

assert prepared.file_object is file_handle
assert prepared.filename == "from_open.bin"
assert prepared.should_close is False
assert prepared.file_object.read() == b"content from open file"


async def test_prepare_upload_with_none() -> None:
"""Test preparing upload with None content (async)."""
prepared = await FileHandlerBase.prepare_upload(content=None)
Expand Down Expand Up @@ -157,6 +186,26 @@ def test_prepare_upload_sync_with_binary_io() -> None:
assert prepared.should_close is False


def test_prepare_upload_sync_with_open_file() -> None:
"""Test preparing upload with a real file handle (not BytesIO) — sync.

Locks in BinaryIO branch behaviour for callers passing the result of opening a real
file, not just ``BytesIO``. Guards against narrowing the dispatch to a too-specific
type.
"""
with tempfile.NamedTemporaryFile(suffix=".txt") as tmp:
tmp.write(b"content from open file")
tmp.flush()

with _open_binary_handle(tmp.name) as file_handle:
prepared = FileHandlerBase.prepare_upload_sync(content=file_handle, name="from_open.bin")

assert prepared.file_object is file_handle
assert prepared.filename == "from_open.bin"
assert prepared.should_close is False
assert prepared.file_object.read() == b"content from open file"


def test_prepare_upload_sync_with_none() -> None:
"""Test preparing upload with None content (sync)."""
prepared = FileHandlerBase.prepare_upload_sync(content=None)
Expand Down
Loading