Merge pull request #449 from jlowin/prune_titles

src/fastmcp/utilities/json_schema.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Mapping, Sequence
 
 
 def _prune_param(schema: dict, param: str) -> dict:
@@ -25,32 +24,55 @@ def _prune_param(schema: dict, param: str) -> dict:
     return schema
 
 
-def _prune_unused_defs(schema: dict) -> dict:
-    """Remove unused definitions from the schema."""
-    # collect all remaining local $ref targets
+def _walk_and_prune(
+    schema: dict,
+    prune_defs: bool = False,
+    prune_titles: bool = False,
+    prune_additional_properties: bool = False,
+) -> dict:
+    """Walk the schema and optionally prune titles, unused definitions, and additionalProperties: false."""
+
+    # Will only be used if prune_defs is True
     used_defs: set[str] = set()
 
-    def walk(node: object) -> None:  # depth-first traversal
-        if isinstance(node, Mapping):
-            ref = node.get("$ref")
-            if isinstance(ref, str) and ref.startswith("#/$defs/"):
-                used_defs.add(ref.split("/")[-1])
+    def walk(node: object) -> None:
+        if isinstance(node, dict):
+            # Process $ref for definition tracking
+            if prune_defs:
+                ref = node.get("$ref")
+                if isinstance(ref, str) and ref.startswith("#/$defs/"):
+                    used_defs.add(ref.split("/")[-1])
+
+            # Remove title if requested
+            if prune_titles and "title" in node:
+                node.pop("title")
+
+            # Remove additionalProperties: false at any level if requested
+            if (
+                prune_additional_properties
+                and node.get("additionalProperties", None) is False
+            ):
+                node.pop("additionalProperties")
+
+            # Walk children
             for v in node.values():
                 walk(v)
-        elif isinstance(node, Sequence) and not isinstance(node, str | bytes):
+
+        elif isinstance(node, list):
             for v in node:
                 walk(v)
 
+    # Traverse the schema once
     walk(schema)
 
-    # remove orphaned definitions
-
-    defs = schema.get("$defs", {})
-    for def_name in list(defs):
-        if def_name not in used_defs:
-            defs.pop(def_name)
-    if not defs:
-        schema.pop("$defs", None)
+    # Remove orphaned definitions if requested
+    if prune_defs:
+        defs = schema.get("$defs", {})
+        for def_name in list(defs):
+            if def_name not in used_defs:
+                defs.pop(def_name)
+        if not defs:
+            schema.pop("$defs", None)
 
     return schema
 
@@ -67,16 +89,32 @@ def compress_schema(
     prune_params: list[str] | None = None,
     prune_defs: bool = True,
     prune_additional_properties: bool = True,
+    prune_titles: bool = False,
 ) -> dict:
     """
     Remove the given parameters from the schema.
 
+    Args:
+        schema: The schema to compress
+        prune_params: List of parameter names to remove from properties
+        prune_defs: Whether to remove unused definitions
+        prune_additional_properties: Whether to remove additionalProperties: false
+        prune_titles: Whether to remove title fields from the schema
     """
+    # Make a copy so we don't modify the original
     schema = copy.deepcopy(schema)
+
+    # Remove specific parameters if requested
     for param in prune_params or []:
         schema = _prune_param(schema, param=param)
-    if prune_defs:
-        schema = _prune_unused_defs(schema)
-    if prune_additional_properties:
-        schema = _prune_additional_properties(schema)
+
+    # Do a single walk to handle pruning operations
+    if prune_defs or prune_titles or prune_additional_properties:
+        schema = _walk_and_prune(
+            schema,
+            prune_defs=prune_defs,
+            prune_titles=prune_titles,
+            prune_additional_properties=prune_additional_properties,
+        )
+
     return schema

tests/utilities/test_json_schema.py

@@ -1,11 +1,21 @@
 from fastmcp.utilities.json_schema import (
-    _prune_additional_properties,
     _prune_param,
-    _prune_unused_defs,
+    _walk_and_prune,
     compress_schema,
 )
 
 
+# Create wrappers for backward compatibility with tests
+def _prune_unused_defs(schema):
+    """Wrapper for _walk_and_prune that only prunes definitions."""
+    return _walk_and_prune(schema, prune_defs=True)
+
+
+def _prune_additional_properties(schema):
+    """Wrapper for _walk_and_prune that only prunes additionalProperties: false."""
+    return _walk_and_prune(schema, prune_additional_properties=True)
+
+
 class TestPruneParam:
     """Tests for the _prune_param function."""
 
@@ -244,3 +254,51 @@ class TestCompressSchema:
         assert "$defs" not in result  # Both defs should be gone
         # Check that additionalProperties was removed
         assert "additionalProperties" not in result
+
+    def test_prune_titles(self):
+        """Test pruning title fields."""
+        schema = {
+            "title": "Root Schema",
+            "type": "object",
+            "properties": {
+                "foo": {"title": "Foo Property", "type": "string"},
+                "bar": {
+                    "title": "Bar Property",
+                    "type": "object",
+                    "properties": {
+                        "nested": {"title": "Nested Property", "type": "string"}
+                    },
+                },
+            },
+        }
+        result = compress_schema(schema, prune_titles=True)
+        assert "title" not in result
+        assert "title" not in result["properties"]["foo"]
+        assert "title" not in result["properties"]["bar"]
+        assert "title" not in result["properties"]["bar"]["properties"]["nested"]
+
+    def test_prune_nested_additional_properties(self):
+        """Test pruning additionalProperties: false at all levels."""
+        schema = {
+            "type": "object",
+            "additionalProperties": False,
+            "properties": {
+                "foo": {
+                    "type": "object",
+                    "additionalProperties": False,
+                    "properties": {
+                        "nested": {
+                            "type": "object",
+                            "additionalProperties": False,
+                        }
+                    },
+                },
+            },
+        }
+        result = compress_schema(schema)
+        assert "additionalProperties" not in result
+        assert "additionalProperties" not in result["properties"]["foo"]
+        assert (
+            "additionalProperties"
+            not in result["properties"]["foo"]["properties"]["nested"]
+        )