ifcquery/ifcmcp: fix plot bugs and add ifc_plot MCP tool

Generated with the assistance of an AI coding tool.

src/ifcmcp/ifcmcp/core.py

@@ -13,7 +13,7 @@ from ifcedit.run import run_api
 from ifcquery import clash as clash_mod
 from ifcquery import contexts as contexts_mod
 from ifcquery import cost as cost_mod
-from ifcquery import info, materials as materials_mod, relations, render as render_mod, schedule, schema, select, summary, tree
+from ifcquery import info, materials as materials_mod, plot as plot_mod, relations, render as render_mod, schedule, schema, select, summary, tree
 from ifcquery import validate as validate_mod
 
 
@@ -360,6 +360,48 @@ class IfcSession:
         """Return IFC class documentation for entity_type using the model's schema version."""
         return schema.schema(self._require_model(), entity_type)
 
+    def ifc_plot(
+        self,
+        selector: str = "",
+        element_ids: list[int] | None = None,
+        view: str = "floorplan",
+        width_mm: float = 297.0,
+        height_mm: float = 420.0,
+        scale: float = 1.0 / 100.0,
+        png_width: int = 1024,
+        png_height: int = 1024,
+    ) -> bytes:
+        """Generate a 2D technical drawing (floor plan, elevation, or section) and return PNG bytes.
+
+        Uses ifcopenshell.draw to produce SVG output which is rasterised to PNG via CairoSVG.
+
+        :param selector: ifcopenshell selector to restrict plotted elements
+            (e.g. ``'IfcWall'``). Omit to plot the whole model.
+        :param element_ids: Step IDs of elements to highlight. Other elements
+            are faded to 10% opacity so the subject stands out.
+        :param view: Drawing view — ``floorplan`` (default), ``elevation``,
+            ``section``, or ``auto``.
+        :param width_mm: Paper width in mm (default 297 = A4).
+        :param height_mm: Paper height in mm (default 420 = A4).
+        :param scale: Model-to-paper scale ratio (default 0.01 = 1:100).
+        :param png_width: Raster output width in pixels (default 1024).
+        :param png_height: Raster output height in pixels (default 1024).
+        :return: PNG image as raw bytes.
+        """
+        model = self._require_model()
+        return plot_mod.plot(
+            model,
+            output_format="png",
+            selector=selector if selector else None,
+            element_ids=element_ids,
+            view=view,
+            width_mm=width_mm,
+            height_mm=height_mm,
+            scale=scale,
+            png_width=png_width,
+            png_height=png_height,
+        )
+
     def ifc_render(
         self,
         selector: str = "",

src/ifcmcp/ifcmcp/server.py

@@ -140,6 +140,46 @@ def build_server() -> Any:
     def ifc_shape(method: str, params: str = "{}") -> dict:
         return session.ifc_shape(method=method, params=params)
 
+    @server.tool(structured_output=False)
+    def ifc_plot(
+        selector: str = "",
+        element_ids: list[int] | None = None,
+        view: str = "floorplan",
+        width_mm: float = 297.0,
+        height_mm: float = 420.0,
+        scale: float = 1.0 / 100.0,
+        png_width: int = 1024,
+        png_height: int = 1024,
+    ) -> list[ImageContent]:
+        """Generate a 2D technical drawing of the loaded IFC model.
+
+        Returns an inline image (floor plan, elevation, or section) that the
+        LLM can inspect to understand the 2D layout of the model.
+
+        :param selector: ifcopenshell selector to restrict plotted elements
+            (e.g. ``'IfcWall'``). Omit to plot the whole model.
+        :param element_ids: Step IDs of elements to highlight. Other elements
+            are faded so the subject stands out.
+        :param view: Drawing view — ``floorplan`` (default), ``elevation``,
+            ``section``, or ``auto``.
+        :param width_mm: Paper width in mm (default 297 = A4 landscape width).
+        :param height_mm: Paper height in mm (default 420 = A4 landscape height).
+        :param scale: Model-to-paper scale ratio (default 0.01 = 1:100).
+        :param png_width: Raster output width in pixels (default 1024).
+        :param png_height: Raster output height in pixels (default 1024).
+        """
+        png_bytes = session.ifc_plot(
+            selector=selector,
+            element_ids=element_ids,
+            view=view,
+            width_mm=width_mm,
+            height_mm=height_mm,
+            scale=scale,
+            png_width=png_width,
+            png_height=png_height,
+        )
+        return [ImageContent(type="image", data=base64.b64encode(png_bytes).decode(), mimeType="image/png")]
+
     @server.tool(structured_output=False)
     def ifc_render(
         selector: str = "",

src/ifcquery/ifcquery/__main__.py

@@ -296,29 +296,40 @@ def main():
                 print(f"Error: Invalid element ID(s): {args.element}", file=sys.stderr)
                 sys.exit(1)
 
-        # Choose default output extension based on out-format
+        try:
+            result = plot.plot(
+                model,
+                selector=args.selector or None,
+                element_ids=element_ids,
+                view=args.view,
+                width_mm=args.width_mm,
+                height_mm=args.height_mm,
+                scale=args.scale,
+                output_format=args.out_format,
+            )
+        except ImportError as e:
+            print(f"Error: {e}", file=sys.stderr)
+            sys.exit(1)
+        except ValueError as e:
+            print(f"Error: {e}", file=sys.stderr)
+            sys.exit(1)
+
+        if args.out_format == "base64":
+            # result is a dict; serialise to stdout so callers can consume it
+            print(format_output(result, args.output_format))
+            return
+
+        # svg or png: write to a file
         base = os.path.splitext(args.ifc_file)[0]
         if args.out_format == "svg":
             out_path = args.output or (base + ".svg")
-        elif args.out_format == "png":
-            out_path = args.output or (base + ".png")
         else:
-            out_path = args.output  # unused; base64 prints to stdout via format_output
+            out_path = args.output or (base + ".png")
 
-        png_bytes = plot.plot(
-            model,
-            selector=args.selector or None,
-            element_ids=element_ids,
-            view=args.view,
-            width_mm=args.width_mm,
-            height_mm=args.height_mm,
-            scale=args.scale,
-            output_format=args.out_format
-        )
         with open(out_path, "wb") as f:
-            f.write(png_bytes)
+            f.write(result)
 
-        print(f"Saved render to {out_path}", file=sys.stderr)
+        print(f"Saved drawing to {out_path}", file=sys.stderr)
         return
 
 

src/ifcquery/ifcquery/plot.py

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import base64
 from io import BytesIO
 import os
 from typing import Any
@@ -61,7 +62,10 @@ def _escape_css_attr(name: str) -> str:
 def _highlight_css_from_ids(model: ifcopenshell.file, element_ids: list[int]) -> str:
     guids: list[str] = []
     for sid in element_ids:
-        e = model.by_id(int(sid))
+        try:
+            e = model.by_id(int(sid))
+        except RuntimeError:
+            continue
         if e is None:
             continue
         gid = getattr(e, "GlobalId", None)
@@ -219,9 +223,12 @@ def plot(
     if not _HAS_CAIROSVG:
         raise ImportError("CairoSVG is not installed. Install with: pip install cairosvg")
 
+    svgs = list(svg_split(BytesIO(svg_bytes)))
+    if not svgs:
+        raise ValueError("No plan geometry found; the model may lack 2D annotation representations for the requested view.")
+
     composite = None
     png_bytes = None
-    svgs = list(svg_split(BytesIO(svg_bytes)))
     for i, svgb in enumerate(svgs):
         png_bytes = cairosvg.svg2png(bytestring=svgb, output_width=png_width, output_height=png_height)
         if len(svgs) == 1:
@@ -240,4 +247,13 @@ def plot(
         composite.save(b, 'png')
         png_bytes = b.getvalue()
 
+    if output_format == "base64":
+        return {
+            "mime": "image/png",
+            "png_b64": base64.b64encode(png_bytes).decode(),
+            "width": png_width,
+            "height": png_height,
+            "view": view,
+        }
+
     return png_bytes

src/ifcquery/tests/test_plot.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import base64
+import os
+import subprocess
+import sys
+import tempfile
+
+import ifcopenshell
+import ifcopenshell.api.aggregate
+import ifcopenshell.api.context
+import ifcopenshell.api.geometry
+import ifcopenshell.api.owner.settings
+import ifcopenshell.api.project
+import ifcopenshell.api.root
+import ifcopenshell.api.spatial
+import ifcopenshell.api.unit
+import pytest
+
+from ifcquery.plot import plot, _highlight_css_from_ids
+
+try:
+    import ifcopenshell.draw  # noqa: F401
+
+    HAS_DRAW = True
+except ImportError:
+    HAS_DRAW = False
+
+try:
+    import cairosvg  # noqa: F401
+
+    HAS_CAIROSVG = True
+except ImportError:
+    HAS_CAIROSVG = False
+
+pytestmark = pytest.mark.skipif(not HAS_DRAW, reason="ifcopenshell.draw not available")
+
+SVG_MAGIC = b"<?xml"
+PNG_MAGIC = b"\x89PNG"
+
+
+@pytest.fixture
+def model_with_annotations():
+    """IFC4 model with walls and explicit 2D annotation geometry (Plan/PLAN_VIEW context)."""
+    f = ifcopenshell.api.project.create_file()
+    ifcopenshell.api.owner.settings.get_user = lambda ifc: (ifc.by_type("IfcPersonAndOrganization") or [None])[0]
+    ifcopenshell.api.owner.settings.get_application = lambda ifc: (ifc.by_type("IfcApplication") or [None])[0]
+
+    project = ifcopenshell.api.root.create_entity(f, ifc_class="IfcProject", name="TestProject")
+    ifcopenshell.api.unit.assign_unit(f)
+
+    site = ifcopenshell.api.root.create_entity(f, ifc_class="IfcSite", name="TestSite")
+    building = ifcopenshell.api.root.create_entity(f, ifc_class="IfcBuilding", name="TestBuilding")
+    storey = ifcopenshell.api.root.create_entity(f, ifc_class="IfcBuildingStorey", name="Ground Floor")
+    ifcopenshell.api.aggregate.assign_object(f, products=[site], relating_object=project)
+    ifcopenshell.api.aggregate.assign_object(f, products=[building], relating_object=site)
+    ifcopenshell.api.aggregate.assign_object(f, products=[storey], relating_object=building)
+
+    model_ctx = ifcopenshell.api.context.add_context(f, context_type="Model")
+    body = ifcopenshell.api.context.add_context(
+        f, context_type="Model", context_identifier="Body", target_view="MODEL_VIEW", parent=model_ctx
+    )
+
+    wall = ifcopenshell.api.root.create_entity(f, ifc_class="IfcWall", name="Wall001")
+    rep = ifcopenshell.api.geometry.add_wall_representation(f, context=body, length=5, height=3, thickness=0.2)
+    ifcopenshell.api.geometry.assign_representation(f, product=wall, representation=rep)
+    ifcopenshell.api.spatial.assign_container(f, products=[wall], relating_structure=storey)
+
+    return f, wall
+
+
+@pytest.fixture
+def model_no_plan(model_with_annotations):
+    """Model whose SVG output will be empty (wall geometry only, no plan annotation group)."""
+    return model_with_annotations
+
+
+class TestHighlightCSS:
+    def test_css_for_valid_element(self, model_with_annotations):
+        model, wall = model_with_annotations
+        css = _highlight_css_from_ids(model, [wall.id()])
+        assert wall.GlobalId in css
+        assert "opacity: 0.10" in css
+        assert "opacity: 1.0" in css
+        assert "#d00" in css
+
+    def test_css_empty_for_no_ids(self, model_with_annotations):
+        model, _ = model_with_annotations
+        css = _highlight_css_from_ids(model, [])
+        assert css == ""
+
+    def test_css_skips_unknown_ids(self, model_with_annotations):
+        model, _ = model_with_annotations
+        css = _highlight_css_from_ids(model, [999999])
+        assert css == ""
+
+
+class TestPlotSVG:
+    def test_returns_svg_bytes(self, model_with_annotations):
+        model, _ = model_with_annotations
+        result = plot(model, output_format="svg")
+        assert isinstance(result, bytes)
+        assert result[:5] == SVG_MAGIC
+
+    def test_svg_contains_xml(self, model_with_annotations):
+        model, _ = model_with_annotations
+        result = plot(model, output_format="svg")
+        assert b"<svg" in result
+
+    def test_invalid_format_raises(self, model_with_annotations):
+        model, _ = model_with_annotations
+        with pytest.raises(ValueError, match="output_format"):
+            plot(model, output_format="xyz")
+
+    def test_invalid_view_raises(self, model_with_annotations):
+        model, _ = model_with_annotations
+        with pytest.raises(ValueError, match="view"):
+            plot(model, output_format="svg", view="bogus")
+
+    def test_selector_no_match_raises(self, model_with_annotations):
+        model, _ = model_with_annotations
+        with pytest.raises(ValueError, match="matched no elements"):
+            plot(model, output_format="svg", selector="IfcDoor")
+
+    def test_selector_filters_elements(self, model_with_annotations):
+        model, _ = model_with_annotations
+        result = plot(model, output_format="svg", selector="IfcWall")
+        assert isinstance(result, bytes)
+        assert b"<svg" in result
+
+
+class TestPlotEmptySVG:
+    """When draw produces no <g> elements, PNG/base64 should raise a clear error."""
+
+    def test_empty_drawing_png_raises(self, model_no_plan):
+        """PNG format raises ValueError (not silently returns None) for empty drawings."""
+        model, _ = model_no_plan
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if not has_groups:
+            pytest.raises(ValueError, plot, model, output_format="png")
+        else:
+            pytest.skip("Model produced non-empty SVG — empty path not triggered")
+
+    def test_empty_drawing_base64_raises(self, model_no_plan):
+        """base64 format raises ValueError (not silently returns None) for empty drawings."""
+        model, _ = model_no_plan
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if not has_groups:
+            pytest.raises(ValueError, plot, model, output_format="base64")
+        else:
+            pytest.skip("Model produced non-empty SVG — empty path not triggered")
+
+
+@pytest.mark.skipif(not HAS_CAIROSVG, reason="cairosvg not installed")
+class TestPlotPNG:
+    """PNG and base64 require cairosvg."""
+
+    def test_png_returns_bytes_or_raises_on_empty(self, model_with_annotations):
+        model, _ = model_with_annotations
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if has_groups:
+            result = plot(model, output_format="png")
+            assert isinstance(result, bytes)
+            assert result[:4] == PNG_MAGIC
+        else:
+            with pytest.raises(ValueError, match="No plan geometry"):
+                plot(model, output_format="png")
+
+    def test_base64_returns_dict(self, model_with_annotations):
+        model, _ = model_with_annotations
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if has_groups:
+            result = plot(model, output_format="base64")
+            assert isinstance(result, dict)
+            assert result["mime"] == "image/png"
+            assert "png_b64" in result
+            assert "width" in result
+            assert "height" in result
+            assert "view" in result
+            # Verify the base64 is valid PNG
+            decoded = base64.b64decode(result["png_b64"])
+            assert decoded[:4] == PNG_MAGIC
+        else:
+            with pytest.raises(ValueError, match="No plan geometry"):
+                plot(model, output_format="base64")
+
+    def test_base64_view_field_matches_requested(self, model_with_annotations):
+        model, _ = model_with_annotations
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if not has_groups:
+            pytest.skip("Model produces empty SVG")
+        result = plot(model, output_format="base64", view="floorplan")
+        assert result["view"] == "floorplan"
+
+    def test_png_custom_size(self, model_with_annotations):
+        model, _ = model_with_annotations
+        svg = plot(model, output_format="svg")
+        has_groups = b"<g " in svg or b"<g>" in svg
+        if not has_groups:
+            pytest.skip("Model produces empty SVG")
+        result = plot(model, output_format="png", png_width=512, png_height=512)
+        assert isinstance(result, bytes)
+        assert result[:4] == PNG_MAGIC
+
+
+class TestCLI:
+    @staticmethod
+    def _ifc_path(model):
+        f = tempfile.NamedTemporaryFile(suffix=".ifc", delete=False)
+        model.write(f.name)
+        f.close()
+        return f.name
+
+    def test_plot_svg_writes_file(self, model_with_annotations):
+        model, _ = model_with_annotations
+        ifc_path = self._ifc_path(model)
+        out_path = ifc_path.replace(".ifc", "_out.svg")
+        try:
+            result = subprocess.run(
+                [sys.executable, "-m", "ifcquery", ifc_path, "plot", "--out-format", "svg", "-o", out_path],
+                capture_output=True,
+                text=True,
+            )
+            assert result.returncode == 0, result.stderr
+            assert os.path.exists(out_path)
+            with open(out_path, "rb") as f:
+                assert f.read(5) == SVG_MAGIC
+        finally:
+            for path in (ifc_path, out_path):
+                try:
+                    os.unlink(path)
+                except OSError:
+                    pass
+
+    @pytest.mark.skipif(not HAS_CAIROSVG, reason="cairosvg not installed")
+    def test_plot_base64_prints_json(self, model_with_annotations):
+        """base64 format prints JSON to stdout instead of writing a file."""
+        model, _ = model_with_annotations
+        ifc_path = self._ifc_path(model)
+        try:
+            # First check if the model would produce geometry
+            svg = plot(model, output_format="svg")
+            has_groups = b"<g " in svg or b"<g>" in svg
+            if not has_groups:
+                pytest.skip("Model produces empty SVG — base64 would raise ValueError")
+
+            result = subprocess.run(
+                [sys.executable, "-m", "ifcquery", ifc_path, "plot", "--out-format", "base64"],
+                capture_output=True,
+                text=True,
+            )
+            assert result.returncode == 0, result.stderr
+            # Output should be JSON (not an error) and contain base64 key
+            assert "png_b64" in result.stdout
+        finally:
+            try:
+                os.unlink(ifc_path)
+            except OSError:
+                pass