diff --git a/packages/capabilities/image-generation/pyproject.toml b/packages/capabilities/image-generation/pyproject.toml index d3b01b70..e0b5fa35 100644 --- a/packages/capabilities/image-generation/pyproject.toml +++ b/packages/capabilities/image-generation/pyproject.toml @@ -27,6 +27,7 @@ Issues = "https://github.com/withceleste/celeste-python/issues" [tool.uv.sources] celeste-ai = { workspace = true } +celeste-bfl = { workspace = true } celeste-google = { workspace = true } [project.entry-points."celeste.packages"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/io.py b/packages/capabilities/image-generation/src/celeste_image_generation/io.py index 40a3afb3..6f724410 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/io.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/io.py @@ -32,6 +32,8 @@ class ImageGenerationUsage(Usage): reasoning_tokens: int | None = None num_images: int | None = None billed_units: float | None = None + input_mp: float | None = None + output_mp: float | None = None class ImageGenerationOutput(Output[ImageArtifact | list[ImageArtifact]]): diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/client.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/client.py index cbdbcab2..45aa3672 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/client.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/client.py @@ -1,14 +1,11 @@ -"""BFL (Black Forest Labs) client implementation for FLUX.2 image generation.""" +"""BFL client implementation for image generation.""" -import asyncio -import json -import time from typing import Any, Unpack -import httpx +from celeste_bfl.images.client import BFLImagesClient from celeste.artifacts import ImageArtifact -from celeste.mime_types import ApplicationMimeType +from celeste.exceptions import ValidationError from celeste.parameters import ParameterMapper from celeste_image_generation.client import ImageGenerationClient from celeste_image_generation.io import ( @@ -18,12 +15,11 @@ ) from celeste_image_generation.parameters import ImageGenerationParameters -from . import config from .parameters import BFL_PARAMETER_MAPPERS -class BFLImageGenerationClient(ImageGenerationClient): - """Black Forest Labs client for image generation.""" +class BFLImageGenerationClient(BFLImagesClient, ImageGenerationClient): + """BFL client for image generation.""" @classmethod def parameter_mappers(cls) -> list[ParameterMapper]: @@ -37,12 +33,8 @@ def _init_request(self, inputs: ImageGenerationInput) -> dict[str, Any]: def _parse_usage(self, response_data: dict[str, Any]) -> ImageGenerationUsage: """Parse usage from response.""" - submit_metadata = response_data.get("_submit_metadata", {}) - cost = submit_metadata.get("cost") - - return ImageGenerationUsage( - billed_units=float(cost) if cost is not None else None, - ) + usage = super()._parse_usage(response_data) + return ImageGenerationUsage(**usage) def _parse_content( self, @@ -55,13 +47,13 @@ def _parse_content( if not sample_url: msg = f"No image URL in {self.provider} response" - raise ValueError(msg) + raise ValidationError(msg) return ImageArtifact(url=sample_url) def _parse_finish_reason( self, response_data: dict[str, Any] - ) -> ImageGenerationFinishReason | None: + ) -> ImageGenerationFinishReason: """Parse finish reason from response.""" status = response_data.get("status") if status == "Ready": @@ -69,81 +61,7 @@ def _parse_finish_reason( elif status in ("Error", "Failed"): error_msg = response_data.get("error", "Generation failed") return ImageGenerationFinishReason(reason="ERROR", message=error_msg) - return None - - async def _make_request( - self, - request_body: dict[str, Any], - **parameters: Unpack[ImageGenerationParameters], - ) -> httpx.Response: - """Make HTTP request(s) and return response object.""" - headers = { - **self.auth.get_headers(), - "Content-Type": ApplicationMimeType.JSON, - "Accept": ApplicationMimeType.JSON, - } - - endpoint = config.ENDPOINT.format(model_id=self.model.id) - - submit_response = await self.http_client.post( - f"{config.BASE_URL}{endpoint}", - headers=headers, - json_body=request_body, - ) - - if submit_response.status_code != 200: - return submit_response - - submit_data = submit_response.json() - polling_url = submit_data.get("polling_url") - - if not polling_url: - msg = f"No polling_url in {self.provider} response" - raise ValueError(msg) - - start_time = time.monotonic() - poll_headers = { - **self.auth.get_headers(), - "Accept": ApplicationMimeType.JSON, - } - - while True: - elapsed = time.monotonic() - start_time - if elapsed >= config.POLLING_TIMEOUT: - msg = f"{self.provider} polling timed out after {config.POLLING_TIMEOUT} seconds" - raise TimeoutError(msg) - - poll_response = await self.http_client.get( - polling_url, - headers=poll_headers, - ) - - if poll_response.status_code != 200: - return poll_response - - poll_data = poll_response.json() - status = poll_data.get("status") - - if status == "Ready": - final_data = { - **poll_data, - "_submit_metadata": submit_data, - } - return httpx.Response( - status_code=200, - content=json.dumps(final_data).encode("utf-8"), - headers={"content-type": "application/json"}, - request=httpx.Request("GET", polling_url), - ) - elif status in ("Error", "Failed"): - return httpx.Response( - status_code=400, - content=json.dumps(poll_data).encode("utf-8"), - headers={"content-type": "application/json"}, - request=httpx.Request("GET", polling_url), - ) - - await asyncio.sleep(config.POLLING_INTERVAL) + return ImageGenerationFinishReason(reason=None) __all__ = ["BFLImageGenerationClient"] diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/config.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/config.py deleted file mode 100644 index 86d7bbf3..00000000 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/config.py +++ /dev/null @@ -1,13 +0,0 @@ -"""BFL (Black Forest Labs) provider configuration for image generation.""" - -# HTTP Configuration -BASE_URL = "https://api.bfl.ai" -ENDPOINT = "/v1/{model_id}" - -# Authentication -AUTH_HEADER_NAME = "x-key" -AUTH_HEADER_PREFIX = "" # Direct API key, no prefix - -# Polling Configuration -POLLING_INTERVAL = 0.5 # seconds between polling attempts -POLLING_TIMEOUT = 120.0 # 2 minutes for image generation diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/parameters.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/parameters.py index 64781148..7a3a575f 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/parameters.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/bfl/parameters.py @@ -1,14 +1,47 @@ -"""BFL parameter mappers for image generation.""" +"""BFL Images parameter mappers for image generation. + +AspectRatioMapper is defined locally (handles "WxH" → width/height transformation). +Other mappers subclass from provider and add capability-specific `name` attribute. +""" from typing import Any -from celeste import Model +from celeste_bfl.images.parameters import ( + GuidanceMapper as _GuidanceMapper, +) +from celeste_bfl.images.parameters import ( + HeightMapper as _HeightMapper, +) +from celeste_bfl.images.parameters import ( + OutputFormatMapper as _OutputFormatMapper, +) +from celeste_bfl.images.parameters import ( + PromptUpsamplingMapper as _PromptUpsamplingMapper, +) +from celeste_bfl.images.parameters import ( + SafetyToleranceMapper as _SafetyToleranceMapper, +) +from celeste_bfl.images.parameters import ( + SeedMapper as _SeedMapper, +) +from celeste_bfl.images.parameters import ( + StepsMapper as _StepsMapper, +) +from celeste_bfl.images.parameters import ( + WidthMapper as _WidthMapper, +) + +from celeste.models import Model from celeste.parameters import ParameterMapper from celeste_image_generation.parameters import ImageGenerationParameter class AspectRatioMapper(ParameterMapper): - """Map aspect_ratio to BFL width/height parameters.""" + """Map aspect_ratio to BFL width/height parameters. + + Converts 'WxH' string to width/height, rounded to nearest multiple of 16. + Delegates to provider's WidthMapper and HeightMapper for the actual mapping. + """ name = ImageGenerationParameter.ASPECT_RATIO @@ -31,130 +64,35 @@ def map( width = ((width + 8) // 16) * 16 height = ((height + 8) // 16) * 16 - request["width"] = width - request["height"] = height + # Delegate to provider mappers + request = _WidthMapper().map(request, width, model) + request = _HeightMapper().map(request, height, model) return request -class PromptUpsamplingMapper(ParameterMapper): - """Map prompt_upsampling parameter to BFL request format.""" - +class PromptUpsamplingMapper(_PromptUpsamplingMapper): name = ImageGenerationParameter.PROMPT_UPSAMPLING - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform prompt_upsampling into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["prompt_upsampling"] = validated_value - return request - - -class SeedMapper(ParameterMapper): - """Map seed parameter to BFL request format.""" +class SeedMapper(_SeedMapper): name = ImageGenerationParameter.SEED - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform seed into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["seed"] = validated_value - return request - - -class SafetyToleranceMapper(ParameterMapper): - """Map safety_tolerance parameter to BFL request format.""" +class SafetyToleranceMapper(_SafetyToleranceMapper): name = ImageGenerationParameter.SAFETY_TOLERANCE - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform safety_tolerance into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["safety_tolerance"] = validated_value - return request - - -class OutputFormatMapper(ParameterMapper): - """Map output_format parameter to BFL request format.""" +class OutputFormatMapper(_OutputFormatMapper): name = ImageGenerationParameter.OUTPUT_FORMAT - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform output_format into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["output_format"] = validated_value - return request - - -class StepsMapper(ParameterMapper): - """Map steps parameter to BFL request format (flex only).""" +class StepsMapper(_StepsMapper): name = ImageGenerationParameter.STEPS - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform steps into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["steps"] = validated_value - return request - - -class GuidanceMapper(ParameterMapper): - """Map guidance parameter to BFL request format (flex only).""" +class GuidanceMapper(_GuidanceMapper): name = ImageGenerationParameter.GUIDANCE - def map( - self, - request: dict[str, Any], - value: object, - model: Model, - ) -> dict[str, Any]: - """Transform guidance into provider request.""" - validated_value = self._validate_value(value, model) - if validated_value is None: - return request - - request["guidance"] = validated_value - return request - BFL_PARAMETER_MAPPERS: list[ParameterMapper] = [ AspectRatioMapper(), diff --git a/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py b/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py index 78565b65..ba3575c8 100644 --- a/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py +++ b/packages/capabilities/image-generation/tests/integration_tests/test_image_generation/test_generate.py @@ -35,6 +35,7 @@ async def test_generate(provider: Provider, model: str, parameters: dict) -> Non client = create_client( capability=Capability.IMAGE_GENERATION, provider=provider, + model=model, ) prompt = "A red apple on a white background" diff --git a/packages/providers/bfl/pyproject.toml b/packages/providers/bfl/pyproject.toml new file mode 100644 index 00000000..dbfe2280 --- /dev/null +++ b/packages/providers/bfl/pyproject.toml @@ -0,0 +1,17 @@ +[project] +name = "celeste-bfl" +version = "0.3.0" +description = "BFL (Black Forest Labs) provider package for Celeste AI" +authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] +license = {text = "Apache-2.0"} +requires-python = ">=3.12" + +[tool.uv.sources] +celeste-ai = { workspace = true } + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["src/celeste_bfl"] diff --git a/packages/providers/bfl/src/celeste_bfl/__init__.py b/packages/providers/bfl/src/celeste_bfl/__init__.py new file mode 100644 index 00000000..0cf373b2 --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/__init__.py @@ -0,0 +1 @@ +"""BFL (Black Forest Labs) provider package for Celeste AI.""" diff --git a/packages/providers/bfl/src/celeste_bfl/images/__init__.py b/packages/providers/bfl/src/celeste_bfl/images/__init__.py new file mode 100644 index 00000000..77493658 --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/images/__init__.py @@ -0,0 +1 @@ +"""BFL Images API provider package.""" diff --git a/packages/providers/bfl/src/celeste_bfl/images/client.py b/packages/providers/bfl/src/celeste_bfl/images/client.py new file mode 100644 index 00000000..d24442f7 --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/images/client.py @@ -0,0 +1,150 @@ +"""BFL Images API client with shared implementation.""" + +import asyncio +import time +from typing import Any + +import httpx + +from celeste.core import UsageField +from celeste.io import FinishReason +from celeste.mime_types import ApplicationMimeType + +from . import config + + +class BFLImagesClient: + """Mixin for BFL Images API operations. + + Provides shared implementation: + - _make_request() - HTTP POST with async polling pattern + - _parse_finish_reason() - Map BFL status to FinishReason + + The BFL API uses async polling: + 1. POST to /v1/{model_id} to submit job + 2. Poll GET polling_url until Ready/Failed + 3. Return final response with merged metadata + + Usage: + class BFLImageGenerationClient(BFLImagesClient, ImageGenerationClient): + def _parse_content(self, response_data, **parameters): + result = response_data.get("result", {}) + # Extract image from result["sample"]... + """ + + async def _make_request( + self, + request_body: dict[str, Any], + **parameters: Any, + ) -> httpx.Response: + """Make HTTP request with async polling for BFL image generation. + + Handles the complete async polling workflow: + 1. Submit job to /v1/{model_id} + 2. Poll polling_url until Ready/Failed + 3. Return response with _submit_metadata for usage parsing + """ + auth_headers = self.auth.get_headers() # type: ignore[attr-defined] + headers = { + **auth_headers, + "Content-Type": ApplicationMimeType.JSON, + "Accept": ApplicationMimeType.JSON, + } + + endpoint = config.BFLImagesEndpoint.CREATE_IMAGE.format(model_id=self.model.id) # type: ignore[attr-defined] + + # Phase 1: Submit job + submit_response = await self.http_client.post( # type: ignore[attr-defined] + f"{config.BASE_URL}{endpoint}", + headers=headers, + json_body=request_body, + ) + + if submit_response.status_code != 200: + return submit_response # type: ignore[no-any-return] + + submit_data = submit_response.json() + polling_url = submit_data.get("polling_url") + + if not polling_url: + msg = f"No polling_url in {self.provider} response" # type: ignore[attr-defined] + raise ValueError(msg) + + # Phase 2: Poll for completion + start_time = time.monotonic() + poll_headers = { + **auth_headers, + "Accept": ApplicationMimeType.JSON, + } + + while True: + elapsed = time.monotonic() - start_time + if elapsed >= config.POLLING_TIMEOUT: + msg = f"{self.provider} polling timed out after {config.POLLING_TIMEOUT} seconds" # type: ignore[attr-defined] + raise TimeoutError(msg) + + poll_response = await self.http_client.get( # type: ignore[attr-defined] + polling_url, + headers=poll_headers, + ) + + if poll_response.status_code != 200: + return poll_response # type: ignore[no-any-return] + + poll_data = poll_response.json() + status = poll_data.get("status") + + if status == "Ready": + # Merge submit metadata into final response for usage parsing + final_data = { + **poll_data, + "_submit_metadata": submit_data, + } + return httpx.Response( + status_code=200, + json=final_data, + request=httpx.Request("GET", polling_url), + ) + elif status in ("Error", "Failed"): + return httpx.Response( + status_code=400, + json=poll_data, + request=httpx.Request("GET", polling_url), + ) + + await asyncio.sleep(config.POLLING_INTERVAL) + + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: + """BFL provides status but not structured finish reasons.""" + return FinishReason(reason=None) + + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, float | None]: + """Map BFL usage fields to unified names. + + Shared by client and streaming across all capabilities. + """ + cost = usage_data.get("cost") + input_mp = usage_data.get("input_mp") + output_mp = usage_data.get("output_mp") + return { + UsageField.BILLED_UNITS: float(cost) if cost is not None else None, + UsageField.INPUT_MP: float(input_mp) if input_mp is not None else None, + UsageField.OUTPUT_MP: float(output_mp) if output_mp is not None else None, + } + + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, float | None]: + """Extract usage data from BFL response.""" + submit_metadata = response_data.get("_submit_metadata", {}) + return self.map_usage_fields(submit_metadata) + + def _parse_content(self, response_data: dict[str, Any]) -> Any: + """Parse result from response.""" + result = response_data.get("result", {}) + if not result: + msg = "No result in response" + raise ValueError(msg) + return result + + +__all__ = ["BFLImagesClient"] diff --git a/packages/providers/bfl/src/celeste_bfl/images/config.py b/packages/providers/bfl/src/celeste_bfl/images/config.py new file mode 100644 index 00000000..b9f55cdc --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/images/config.py @@ -0,0 +1,16 @@ +"""Configuration for BFL Images API.""" + +from enum import StrEnum + + +class BFLImagesEndpoint(StrEnum): + """Endpoints for Images API.""" + + CREATE_IMAGE = "/v1/{model_id}" + + +BASE_URL = "https://api.bfl.ai" + +# Polling Configuration +POLLING_INTERVAL = 0.5 # seconds between polling attempts +POLLING_TIMEOUT = 120.0 # 2 minutes timeout diff --git a/packages/providers/bfl/src/celeste_bfl/images/parameters.py b/packages/providers/bfl/src/celeste_bfl/images/parameters.py new file mode 100644 index 00000000..3eacc702 --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/images/parameters.py @@ -0,0 +1,161 @@ +"""BFL Images API parameter mappers.""" + +from typing import Any + +from celeste.constraints import Int +from celeste.models import Model +from celeste.parameters import ParameterMapper + + +class WidthMapper(ParameterMapper): + """Map width to BFL width field.""" + + def map( + self, + request: dict[str, Any], + value: int | str | float | None, + model: Model, + ) -> dict[str, Any]: + """Transform width into provider request.""" + if value is None: + return request + + request["width"] = Int()(value) + return request + + +class HeightMapper(ParameterMapper): + """Map height to BFL height field.""" + + def map( + self, + request: dict[str, Any], + value: int | str | float | None, + model: Model, + ) -> dict[str, Any]: + """Transform height into provider request.""" + if value is None: + return request + + request["height"] = Int()(value) + return request + + +class PromptUpsamplingMapper(ParameterMapper): + """Map prompt_upsampling to BFL prompt_upsampling field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform prompt_upsampling into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["prompt_upsampling"] = validated_value + return request + + +class SeedMapper(ParameterMapper): + """Map seed to BFL seed field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform seed into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["seed"] = validated_value + return request + + +class SafetyToleranceMapper(ParameterMapper): + """Map safety_tolerance to BFL safety_tolerance field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform safety_tolerance into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["safety_tolerance"] = validated_value + return request + + +class OutputFormatMapper(ParameterMapper): + """Map output_format to BFL output_format field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform output_format into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["output_format"] = validated_value + return request + + +class StepsMapper(ParameterMapper): + """Map steps to BFL steps field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform steps into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["steps"] = validated_value + return request + + +class GuidanceMapper(ParameterMapper): + """Map guidance to BFL guidance field.""" + + def map( + self, + request: dict[str, Any], + value: object, + model: Model, + ) -> dict[str, Any]: + """Transform guidance into provider request.""" + validated_value = self._validate_value(value, model) + if validated_value is None: + return request + + request["guidance"] = validated_value + return request + + +__all__ = [ + "GuidanceMapper", + "HeightMapper", + "OutputFormatMapper", + "PromptUpsamplingMapper", + "SafetyToleranceMapper", + "SeedMapper", + "StepsMapper", + "WidthMapper", +] diff --git a/packages/providers/bfl/src/celeste_bfl/images/utils.py b/packages/providers/bfl/src/celeste_bfl/images/utils.py new file mode 100644 index 00000000..f80924e3 --- /dev/null +++ b/packages/providers/bfl/src/celeste_bfl/images/utils.py @@ -0,0 +1,27 @@ +"""BFL Images API utilities.""" + +import base64 + +from celeste.artifacts import ImageArtifact + + +def encode_image(image: ImageArtifact) -> str: + """Encode ImageArtifact to base64 string or return URL. + + BFL API accepts either a URL or base64-encoded image data. + """ + if image.url: + return image.url + elif image.data: + if isinstance(image.data, bytes): + return base64.b64encode(image.data).decode("utf-8") + return image.data + elif image.path: + with open(image.path, "rb") as f: + return base64.b64encode(f.read()).decode("utf-8") + else: + msg = "ImageArtifact must have url, data, or path" + raise ValueError(msg) + + +__all__ = ["encode_image"] diff --git a/packages/providers/bfl/src/celeste_bfl/py.typed b/packages/providers/bfl/src/celeste_bfl/py.typed new file mode 100644 index 00000000..e69de29b diff --git a/src/celeste/constraints.py b/src/celeste/constraints.py index 9c73b990..9ec63878 100644 --- a/src/celeste/constraints.py +++ b/src/celeste/constraints.py @@ -129,10 +129,21 @@ def __call__(self, value: str) -> str: class Int(Constraint): """Integer type constraint.""" - def __call__(self, value: int) -> int: - """Validate value is an integer.""" - # isinstance(True, int) is True, so exclude bools explicitly - if not isinstance(value, int) or isinstance(value, bool): + def __call__(self, value: int | str | float) -> int: + """Validate value is an integer or convert string/float to int.""" + if isinstance(value, str): + if not value.lstrip("-").isdigit(): + msg = f"Must be int, got {value!r}" + raise ConstraintViolationError(msg) + return int(value) + + if isinstance(value, float): + if not value.is_integer(): + msg = f"Must be int, got {value!r}" + raise ConstraintViolationError(msg) + return int(value) + + if not isinstance(value, int): msg = f"Must be int, got {type(value).__name__}" raise ConstraintViolationError(msg) diff --git a/tests/unit_tests/test_constraints.py b/tests/unit_tests/test_constraints.py index 42d1bd05..aeae756a 100644 --- a/tests/unit_tests/test_constraints.py +++ b/tests/unit_tests/test_constraints.py @@ -262,26 +262,37 @@ def test_validates_integer_value(self) -> None: assert result == 42 - def test_rejects_float_value(self) -> None: - """Test that float raises ConstraintViolationError.""" + def test_accepts_boolean_as_int(self) -> None: + """Test that bool is accepted as int (True=1, False=0).""" constraint = Int() + assert constraint(True) == 1 + assert constraint(False) == 0 - with pytest.raises(ConstraintViolationError, match=r"Must be int, got float"): - constraint(42.0) # type: ignore[arg-type] + def test_accepts_valid_string(self) -> None: + """Test that valid integer string is converted.""" + constraint = Int() + assert constraint("42") == 42 + assert constraint("-10") == -10 - def test_rejects_boolean_value(self) -> None: - """Test that bool raises ConstraintViolationError despite isinstance(True, int).""" + def test_rejects_invalid_string(self) -> None: + """Test that non-integer string raises ConstraintViolationError.""" constraint = Int() - with pytest.raises(ConstraintViolationError, match=r"Must be int, got bool"): - constraint(True) + with pytest.raises(ConstraintViolationError, match=r"Must be int, got 'abc'"): + constraint("abc") - def test_rejects_string_value(self) -> None: - """Test that string raises ConstraintViolationError.""" + def test_accepts_whole_float(self) -> None: + """Test that whole number float is converted.""" + constraint = Int() + assert constraint(1.0) == 1 + assert constraint(-5.0) == -5 + + def test_rejects_fractional_float(self) -> None: + """Test that fractional float raises ConstraintViolationError.""" constraint = Int() - with pytest.raises(ConstraintViolationError, match=r"Must be int, got str"): - constraint("42") # type: ignore[arg-type] + with pytest.raises(ConstraintViolationError, match=r"Must be int, got 1.1"): + constraint(1.1) class TestFloat: