From 0e083bf99a90d5433b989512215e352eaaaf8da5 Mon Sep 17 00:00:00 2001 From: Xin Kelly Date: Thu, 12 Mar 2026 16:26:34 +0100 Subject: [PATCH] Add DiagramsAPI.download_converted_file for diagram convert job downloads - Add download_converted_file method to download PNG/SVG for a page of a convert job - Support file_id or file_external_id, page, and mime_type parameters - Add unit tests for the new method - Add diagram convert download to API client retryable endpoints - Add verification script for end-to-end testing Made-with: Cursor --- cognite/client/_api/diagrams.py | 58 ++++++++++++++++++ scripts/verify_diagram_download.py | 70 ++++++++++++++++++++++ tests/tests_unit/test_api/test_diagrams.py | 47 +++++++++++++++ tests/tests_unit/test_api_client.py | 1 + 4 files changed, 176 insertions(+) create mode 100644 scripts/verify_diagram_download.py diff --git a/cognite/client/_api/diagrams.py b/cognite/client/_api/diagrams.py index b9e24f0c95..68042fe911 100644 --- a/cognite/client/_api/diagrams.py +++ b/cognite/client/_api/diagrams.py @@ -394,3 +394,61 @@ def convert(self, detect_job: DiagramDetectResults) -> DiagramConvertResults: items=self._process_detect_job(detect_job), job_cls=DiagramConvertResults, ) + + def download_converted_file( + self, + job_id: int, + *, + file_id: int | None = None, + file_external_id: str | None = None, + page: int = 1, + mime_type: Literal["image/png", "image/svg+xml"] = "image/png", + ) -> bytes: + """Download a converted diagram file (PNG or SVG) for a specific page of a completed convert job. + + The converted file is the output of a diagram convert job, which produces interactive + diagram images with detected annotations highlighted. Use this method to retrieve + the binary content of a specific page. + + Args: + job_id (int): The ID of the completed diagram convert job. + file_id (int | None): The CDF file ID. Either file_id or file_external_id must be provided. + file_external_id (str | None): The CDF file external ID. Either file_id or file_external_id must be provided. + page (int): The page number to download (1-indexed). Defaults to 1. + mime_type (Literal["image/png", "image/svg+xml"]): The desired output format. Defaults to "image/png". + + Returns: + bytes: The binary content of the converted diagram (PNG or SVG). + + Raises: + ValueError: If neither file_id nor file_external_id is provided, or if both are provided. + + Examples: + Download PNG of page 1 from a convert job: + + >>> from cognite.client import CogniteClient + >>> client = CogniteClient() + >>> png_bytes = client.diagrams.download_converted_file( + ... job_id=123, file_id=456, page=1 + ... ) + + Download SVG using file external ID: + + >>> svg_bytes = client.diagrams.download_converted_file( + ... job_id=123, file_external_id="my-diagram.pdf", mime_type="image/svg+xml" + ... ) + """ + if (file_id is None) == (file_external_id is None): + raise ValueError("Exactly one of file_id or file_external_id must be provided") + params: dict[str, Any] = {"page": page} + if file_id is not None: + params["file_id"] = file_id + else: + params["file_external_id"] = file_external_id + res = self._do_request( + "GET", + f"{self._RESOURCE_PATH}/convert/{job_id}/download", + params={to_camel_case(k): v for k, v in params.items()}, + accept=mime_type, + ) + return res.content diff --git a/scripts/verify_diagram_download.py b/scripts/verify_diagram_download.py new file mode 100644 index 0000000000..74a20efb45 --- /dev/null +++ b/scripts/verify_diagram_download.py @@ -0,0 +1,70 @@ +#!/usr/bin/env python3 +"""Verification script for diagram_download_converted_file. + +Run this script to verify the usefulness of the diagram download converted file endpoint. +Requires: CogniteClient configured with valid credentials and a completed diagram convert job. + +Usage: + # Using file_id (from a completed convert job): + python scripts/verify_diagram_download.py --job-id 123 --file-id 456 + + # Using file_external_id: + python scripts/verify_diagram_download.py --job-id 123 --file-external-id "my-diagram.pdf" + + # Download as SVG instead of PNG: + python scripts/verify_diagram_download.py --job-id 123 --file-id 456 --format svg + + # Download a specific page: + python scripts/verify_diagram_download.py --job-id 123 --file-id 456 --page 2 +""" +from __future__ import annotations + +import argparse +from pathlib import Path + +from cognite.client import CogniteClient + + +def main() -> None: + parser = argparse.ArgumentParser(description="Verify diagram_download_converted_file endpoint") + parser.add_argument("--job-id", type=int, required=True, help="Diagram convert job ID") + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--file-id", type=int, help="CDF file ID") + group.add_argument("--file-external-id", type=str, help="CDF file external ID") + parser.add_argument("--page", type=int, default=1, help="Page number (default: 1)") + parser.add_argument( + "--format", + choices=["png", "svg"], + default="png", + help="Output format (default: png)", + ) + parser.add_argument( + "--output", + type=Path, + help="Save to file (optional). If not set, prints size and first bytes.", + ) + args = parser.parse_args() + + client = CogniteClient() + + mime_type = "image/png" if args.format == "png" else "image/svg+xml" + kwargs: dict = {"job_id": args.job_id, "page": args.page, "mime_type": mime_type} + if args.file_id is not None: + kwargs["file_id"] = args.file_id + else: + kwargs["file_external_id"] = args.file_external_id + + print(f"Downloading converted diagram: job_id={args.job_id}, page={args.page}, format={args.format}") + content = client.diagrams.download_converted_file(**kwargs) + + print(f"Downloaded {len(content)} bytes") + if args.output: + args.output.write_bytes(content) + print(f"Saved to {args.output}") + else: + preview = content[:50] if len(content) >= 50 else content + print(f"First bytes (hex): {preview.hex()[:80]}...") + + +if __name__ == "__main__": + main() diff --git a/tests/tests_unit/test_api/test_diagrams.py b/tests/tests_unit/test_api/test_diagrams.py index 9289bfa5d2..f3baeb2cb1 100644 --- a/tests/tests_unit/test_api/test_diagrams.py +++ b/tests/tests_unit/test_api/test_diagrams.py @@ -88,3 +88,50 @@ def test_run_diagram_detect( job_bundle, _unposted_jobs = cognite_client.diagrams.detect( file_ids=file_ids, entities=entities, multiple_jobs=True ) + + +class TestDiagramDownloadConvertedFile: + @patch.object(DiagramsAPI, "_do_request") + def test_download_converted_file_with_file_id( + self, mocked_do_request: MagicMock, cognite_client: CogniteClient + ) -> None: + mock_response = Mock() + mock_response.content = b"fake-png-bytes" + mocked_do_request.return_value = mock_response + + result = cognite_client.diagrams.download_converted_file( + job_id=123, file_id=456, page=2, mime_type="image/png" + ) + + assert result == b"fake-png-bytes" + mocked_do_request.assert_called_once() + call_kwargs = mocked_do_request.call_args[1] + assert call_kwargs["accept"] == "image/png" + assert call_kwargs["params"] == {"fileId": 456, "page": 2} + + @patch.object(DiagramsAPI, "_do_request") + def test_download_converted_file_with_file_external_id( + self, mocked_do_request: MagicMock, cognite_client: CogniteClient + ) -> None: + mock_response = Mock() + mock_response.content = b"fake-svg-bytes" + mocked_do_request.return_value = mock_response + + result = cognite_client.diagrams.download_converted_file( + job_id=789, file_external_id="my-diagram.pdf", mime_type="image/svg+xml" + ) + + assert result == b"fake-svg-bytes" + mocked_do_request.assert_called_once() + call_kwargs = mocked_do_request.call_args[1] + assert call_kwargs["accept"] == "image/svg+xml" + assert call_kwargs["params"] == {"fileExternalId": "my-diagram.pdf", "page": 1} + + def test_download_converted_file_requires_file_identifier(self, cognite_client: CogniteClient) -> None: + with pytest.raises(ValueError, match="Exactly one of file_id or file_external_id must be provided"): + cognite_client.diagrams.download_converted_file(job_id=123) + + with pytest.raises(ValueError, match="Exactly one of file_id or file_external_id must be provided"): + cognite_client.diagrams.download_converted_file( + job_id=123, file_id=456, file_external_id="both-provided" + ) diff --git a/tests/tests_unit/test_api_client.py b/tests/tests_unit/test_api_client.py index 55b21f361c..f7aa4cd196 100644 --- a/tests/tests_unit/test_api_client.py +++ b/tests/tests_unit/test_api_client.py @@ -1368,6 +1368,7 @@ def test_is_retryable_resource_api_endpoints(self, api_client_with_token, method ("POST", "https://api.cognitedata.com/api/v1/projects/bla/context/diagram/convert", True), ("POST", "https://api.cognitedata.com/api/v1/projects/bla/context/diagram/detect", True), ("GET", "https://api.cognitedata.com/api/v1/projects/bla/context/diagram/convert/123", True), + ("GET", "https://api.cognitedata.com/api/v1/projects/bla/context/diagram/convert/123/download", True), ("GET", "https://api.cognitedata.com/api/v1/projects/bla/context/diagram/detect/456", True), # Simulators ("POST", "https://api.cognitedata.com/api/v1/projects/bla/simulators/list", True),