Merge pull request #68 from UiPath/feat/add-schema-helpers

feat: add schema generation helpers

Author: radu-mocanu

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-runtime"
-version = "0.4.0"
+version = "0.4.1"
 description = "Runtime abstractions and interfaces for building agents and automation scripts in the UiPath ecosystem"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath/runtime/schema.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Any
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -68,9 +69,174 @@ class UiPathRuntimeSchema(BaseModel):
     model_config = COMMON_MODEL_SCHEMA
 
 
+def _get_job_attachment_definition() -> dict[str, Any]:
+    """Get the job-attachment definition schema for UiPath attachments.
+
+    Returns:
+        The JSON schema definition for a UiPath job attachment.
+    """
+    return {
+        "type": "object",
+        "required": ["ID"],
+        "x-uipath-resource-kind": "JobAttachment",
+        "properties": {
+            "ID": {"type": "string"},
+            "FullName": {"type": "string"},
+            "MimeType": {"type": "string"},
+            "Metadata": {
+                "type": "object",
+                "additionalProperties": {"type": "string"},
+            },
+        },
+    }
+
+
+def transform_attachments(schema: dict[str, Any]) -> dict[str, Any]:
+    """Transform UiPathAttachment references in a JSON schema to use $ref.
+
+    This function recursively traverses a JSON schema and replaces any objects
+    with title="UiPathAttachment" with a $ref to "#/definitions/job-attachment",
+    adding the job-attachment definition to the schema's definitions section.
+
+    Args:
+        schema: The JSON schema to transform (will not be modified in-place).
+
+    Returns:
+        A new schema with UiPathAttachment references replaced by $ref.
+
+    Example:
+        >>> schema = {
+        ...     "type": "object",
+        ...     "properties": {
+        ...         "file": {
+        ...             "title": "UiPathAttachment",
+        ...             "type": "object",
+        ...             "properties": {...}
+        ...         }
+        ...     }
+        ... }
+        >>> result = transform_attachments(schema)
+        >>> result["properties"]["file"]
+        {"$ref": "#/definitions/job-attachment"}
+        >>> "job-attachment" in result["definitions"]
+        True
+    """
+    result = copy.deepcopy(schema)
+    has_attachments = False
+
+    def transform_recursive(obj: Any) -> Any:
+        """Recursively transform the schema object."""
+        nonlocal has_attachments
+
+        if isinstance(obj, dict):
+            if obj.get("title") == "UiPathAttachment" and obj.get("type") == "object":
+                has_attachments = True
+                return {"$ref": "#/definitions/job-attachment"}
+
+            return {key: transform_recursive(value) for key, value in obj.items()}
+
+        elif isinstance(obj, list):
+            return [transform_recursive(item) for item in obj]
+
+        else:
+            # Return primitive values as-is
+            return obj
+
+    result = transform_recursive(result)
+
+    # add the job-attachment definition if any are present
+    if has_attachments:
+        if "definitions" not in result:
+            result["definitions"] = {}
+        result["definitions"]["job-attachment"] = _get_job_attachment_definition()
+
+    return result
+
+
+def transform_references(schema, root=None, visited=None):
+    """Recursively resolves $ref references in a JSON schema, handling circular references.
+
+    Returns:
+        tuple: (resolved_schema, has_circular_dependency)
+    """
+    if root is None:
+        root = schema
+
+    if visited is None:
+        visited = set()
+
+    has_circular = False
+
+    if isinstance(schema, dict):
+        if "$ref" in schema:
+            ref_path = schema["$ref"]
+
+            if ref_path in visited:
+                # Circular dependency detected
+                return {
+                    "type": "object",
+                    "description": f"Circular reference to {ref_path}",
+                }, True
+
+            visited.add(ref_path)
+
+            # Resolve the reference
+            ref_parts = ref_path.lstrip("#/").split("/")
+            ref_schema = root
+            for part in ref_parts:
+                ref_schema = ref_schema.get(part, {})
+
+            result, circular = transform_references(ref_schema, root, visited)
+            has_circular = has_circular or circular
+
+            # Remove from visited after resolution (allows the same ref in different branches)
+            visited.discard(ref_path)
+
+            return result, has_circular
+
+        resolved_dict = {}
+        for k, v in schema.items():
+            resolved_value, circular = transform_references(v, root, visited)
+            resolved_dict[k] = resolved_value
+            has_circular = has_circular or circular
+        return resolved_dict, has_circular
+
+    elif isinstance(schema, list):
+        resolved_list = []
+        for item in schema:
+            resolved_item, circular = transform_references(item, root, visited)
+            resolved_list.append(resolved_item)
+            has_circular = has_circular or circular
+        return resolved_list, has_circular
+
+    return schema, False
+
+
+def transform_nullable_types(
+    schema: dict[str, Any] | list[Any] | Any,
+) -> dict[str, Any] | list[Any]:
+    """Process the schema to handle nullable types by removing anyOf with null and keeping the base type."""
+    if isinstance(schema, dict):
+        if "anyOf" in schema and len(schema["anyOf"]) == 2:
+            types = [t.get("type") for t in schema["anyOf"]]
+            if "null" in types:
+                non_null_type = next(
+                    t for t in schema["anyOf"] if t.get("type") != "null"
+                )
+                return non_null_type
+
+        return {k: transform_nullable_types(v) for k, v in schema.items()}
+    elif isinstance(schema, list):
+        return [transform_nullable_types(item) for item in schema]
+    return schema
+
+
 __all__ = [
     "UiPathRuntimeSchema",
     "UiPathRuntimeGraph",
     "UiPathRuntimeNode",
     "UiPathRuntimeEdge",
+    "transform_nullable_types",
+    "transform_references",
+    "transform_attachments",
 ]