feat: Added support for commented JSON; simplified instrumentation exports

Makefile

@@ -11,5 +11,5 @@ lint-fix:
 test:
 	uv run pytest
 
-test-coverage:
+test-cov:
 	uv run coverage run -m pytest -q && uv run coverage report --fail-under=90

README.md

@@ -266,13 +266,15 @@ rendered.normalized_lines # List of normalized lines
 
 # Parsing - Full Document
 rendered.as_json()               # Parse as JSON
+rendered.as_json(allow_comments=True)  # Parse JSON with // and /* */ comments
 rendered.as_yaml()               # Parse as YAML (requires pyyaml)
 rendered.as_xml(strict=False)    # Parse as XML (strict=True for single root)
 rendered.as_markdown_sections()  # Parse markdown headings
 rendered.markdown_section("title") # Find markdown section by title
 
 # Parsing - Fenced Code Blocks
 rendered.as_json_blocks()        # Extract all ```json blocks
+rendered.as_json_blocks(allow_comments=True)  # With comment support
 rendered.as_yaml_blocks()        # Extract all ```yaml blocks
 rendered.as_xml_blocks()         # Extract all ```xml blocks
 
@@ -332,32 +334,25 @@ In templates, use comment-based markers to define sections and trace events:
 
 Comment markers are automatically transformed when `test_mode=True`. This allows jinjatest to be a dev-only dependency since the comments are valid Jinja syntax that render as empty strings in production.
 
-#### Using with Any Jinja Environment
-
-You can add instrumentation to any Jinja environment using `instrument()`:
+You can also use a pre-configured Jinja environment with `TemplateSpec`:
 
 ```python
 from jinja2 import Environment, FileSystemLoader
-from jinjatest import TemplateSpec, instrument
+from jinjatest import TemplateSpec
 
-# Patch any existing Jinja environment
+# Use your own Jinja environment
 env = Environment(loader=FileSystemLoader("templates/"))
-instrument(env)  # Adds `jt` global
+env.globals["my_filter"] = lambda x: x.upper()
 
-# Load template with comment markers transformed
+# TemplateSpec handles instrumentation automatically
 spec = TemplateSpec.from_file("my_template.j2", env=env)
 rendered = spec.render({"name": "World"})
 
 # Check traces after rendering
 if rendered.has_trace("some_event"):
     print("Event was triggered")
-
-# For production, use test_mode=False (markers become no-ops)
-instrument(env, test_mode=False)
 ```
 
-This is useful when you want to add instrumentation to an existing Jinja setup.
-
 ## Pytest Integration
 
 jinjatest provides pytest fixtures automatically:

jinjatest/__init__.py

@@ -28,20 +28,13 @@ def test_welcome_pro_user():
 
 from jinjatest.asserts import PromptAssertionError, PromptAsserts, assert_no_undefined
 from jinjatest.instrumentation import (
-    AnchorIndex,
     ProductionInstrumentation,
     TestInstrumentation,
-    TraceRecorder,
-    create_instrumentation,
-    instrument,
 )
 from jinjatest.markers import (
-    MarkerTransform,
     TemplateMarkers,
     discover_markers,
     has_markers,
-    load_template_with_markers,
-    transform_markers,
 )
 from jinjatest.parsers import (
     FencedBlock,
@@ -95,19 +88,12 @@ def test_welcome_pro_user():
     "extract_fenced_blocks",
     "parse_fenced_blocks",
     "FencedBlock",
-    # Instrumentation
-    "instrument",
-    "create_instrumentation",
+    # Instrumentation (for type hints)
     "TestInstrumentation",
     "ProductionInstrumentation",
-    "TraceRecorder",
-    "AnchorIndex",
     # Markers (comment-based)
-    "transform_markers",
     "has_markers",
     "discover_markers",
-    "load_template_with_markers",
-    "MarkerTransform",
     "TemplateMarkers",
     # Utilities
     "normalize_text",

jinjatest/instrumentation.py

@@ -10,10 +10,6 @@
 
 import re
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from jinja2 import Environment
 
 # Sentinel markers for anchors (using ASCII record separator character)
 ANCHOR_START = "\x1e"
@@ -171,44 +167,3 @@ def create_instrumentation(
     if test_mode:
         return TestInstrumentation()
     return ProductionInstrumentation()
-
-
-def instrument(
-    env: Environment,
-    *,
-    test_mode: bool = True,
-    global_name: str = "jt",
-) -> TestInstrumentation | ProductionInstrumentation:
-    """Patch a Jinja environment with instrumentation.
-
-    This function adds instrumentation to any Jinja environment, enabling
-    the use of anchors and traces in templates. Use this when you want to
-    add jinjatest instrumentation to an existing environment without using
-    TemplateSpec.
-
-    Args:
-        env: The Jinja Environment to instrument.
-        test_mode: If True, enable anchors/traces. If False, no-op mode.
-        global_name: Name of the template global variable (default: "jt").
-
-    Returns:
-        The instrumentation instance added to the environment.
-
-    Example:
-        from jinja2 import Environment, FileSystemLoader
-        from jinjatest import instrument
-
-        # Patch any Jinja environment
-        env = Environment(loader=FileSystemLoader("templates/"))
-        inst = instrument(env)
-
-        # Now templates can use {{ jt.anchor("section") }} and {{ jt.trace("event") }}
-        template = env.get_template("my_template.j2")
-        result = template.render({"name": "World"})
-
-        # In production, use test_mode=False for no-op instrumentation
-        instrument(env, test_mode=False)
-    """
-    inst = create_instrumentation(test_mode=test_mode)
-    env.globals[global_name] = inst
-    return inst

jinjatest/markers.py

@@ -158,11 +158,12 @@ def load_template_with_markers(
 
     Example:
         from jinja2 import Environment, FileSystemLoader
-        from jinjatest import instrument
+        from jinjatest.instrumentation import create_instrumentation
         from jinjatest.markers import load_template_with_markers
 
         env = Environment(loader=FileSystemLoader("templates/"))
-        inst = instrument(env)
+        inst = create_instrumentation(test_mode=True)
+        env.globals["jt"] = inst
         template = load_template_with_markers(env, "my_prompt.j2", inst)
         result = template.render({"name": "World"})
     """

jinjatest/parsers/json_parser.py

@@ -14,20 +14,101 @@ def __init__(self, message: str, original_error: Exception | None = None) -> Non
         self.original_error = original_error
 
 
-def parse_json(text: str) -> Any:
+def _is_escaped(text: str, pos: int) -> bool:
+    """Check if character at pos is escaped by counting preceding backslashes."""
+    num_backslashes = 0
+    pos -= 1
+    while pos >= 0 and text[pos] == "\\":
+        num_backslashes += 1
+        pos -= 1
+    # Odd number of backslashes means the character is escaped
+    return num_backslashes % 2 == 1
+
+
+def _strip_json_comments(text: str) -> str:
+    """Strip C-style comments from JSON text.
+
+    Handles:
+    - Single-line comments: // comment
+    - Multi-line comments: /* comment */
+
+    Properly handles comments inside strings (leaves them untouched),
+    including escaped quotes and escaped backslashes.
+    """
+    result: list[str] = []
+    i = 0
+    n = len(text)
+    in_string = False
+
+    while i < n:
+        char = text[i]
+
+        # Handle string state - check for unescaped quotes
+        if char == '"' and not _is_escaped(text, i):
+            in_string = not in_string
+            result.append(char)
+            i += 1
+        elif in_string:
+            # Inside a string - copy everything including escape sequences
+            result.append(char)
+            i += 1
+        elif char == "/" and i + 1 < n:
+            next_char = text[i + 1]
+            if next_char == "/":
+                # Single-line comment - skip to end of line
+                i += 2
+                while i < n and text[i] != "\n":
+                    i += 1
+            elif next_char == "*":
+                # Multi-line comment - skip to */
+                i += 2
+                while i + 1 < n and not (text[i] == "*" and text[i + 1] == "/"):
+                    i += 1
+                i += 2  # Skip the closing */
+            else:
+                result.append(char)
+                i += 1
+        else:
+            result.append(char)
+            i += 1
+
+    return "".join(result)
+
+
+def parse_json(text: str, *, allow_comments: bool = False) -> Any:
     """Parse text as JSON.
 
     Args:
         text: The text to parse as JSON.
+        allow_comments: If True, strip C-style comments (// and /* */)
+            before parsing. Useful for JSONC-style configuration files.
+            Default is False for strict JSON compliance.
 
     Returns:
         The parsed JSON value (dict, list, str, int, float, bool, or None).
 
     Raises:
         JSONParseError: If parsing fails.
+
+    Example:
+        >>> parse_json('{"key": "value"}')
+        {'key': 'value'}
+
+        >>> parse_json('''
+        ... {
+        ...     // This is a comment
+        ...     "key": "value"
+        ... }
+        ... ''', allow_comments=True)
+        {'key': 'value'}
     """
+    text = text.strip()
+
+    if allow_comments:
+        text = _strip_json_comments(text)
+
     try:
-        return json.loads(text.strip())
+        return json.loads(text)
     except json.JSONDecodeError as e:
         raise JSONParseError(
             f"Failed to parse JSON: {e.msg} at line {e.lineno}, column {e.colno}",

jinjatest/rendered.py

@@ -8,6 +8,7 @@
 
 import re
 from dataclasses import dataclass, field
+from functools import partial
 from typing import TYPE_CHECKING, Any
 
 from jinjatest.instrumentation import AnchorIndex
@@ -170,16 +171,21 @@ def has_section(self, name: str) -> bool:
 
     # Parsing methods
 
-    def as_json(self) -> Any:
+    def as_json(self, *, allow_comments: bool = False) -> Any:
         """Parse the rendered text as JSON.
 
+        Args:
+            allow_comments: If True, strip C-style comments (// and /* */)
+                before parsing. Useful for JSONC-style configuration files.
+                Default is False for strict JSON compliance.
+
         Returns:
             The parsed JSON value.
 
         Raises:
             JSONParseError: If parsing fails.
         """
-        return parse_json(self.clean_text)
+        return parse_json(self.clean_text, allow_comments=allow_comments)
 
     def as_yaml(self) -> Any:
         """Parse the rendered text as YAML.
@@ -236,16 +242,22 @@ def as_xml(self, *, strict: bool = False) -> XMLElement | list[XMLElement]:
 
     # Fenced code block parsing methods
 
-    def as_json_blocks(self) -> list[Any]:
+    def as_json_blocks(self, *, allow_comments: bool = False) -> list[Any]:
         """Extract and parse all ```json fenced code blocks.
 
+        Args:
+            allow_comments: If True, strip C-style comments (// and /* */)
+                before parsing. Useful for JSONC-style configuration files.
+                Default is False for strict JSON compliance.
+
         Returns:
             List of parsed JSON objects, one per block found.
 
         Raises:
             JSONParseError: If any block contains invalid JSON.
         """
-        return parse_fenced_blocks(self.clean_text, "json", parse_json)
+        parser = partial(parse_json, allow_comments=allow_comments)
+        return parse_fenced_blocks(self.clean_text, "json", parser)
 
     def as_yaml_blocks(self) -> list[Any]:
         """Extract and parse all ```yaml fenced code blocks.

jinjatest/spec.py

@@ -28,7 +28,6 @@
     ProductionInstrumentation,
     TestInstrumentation,
     create_instrumentation,
-    instrument,
 )
 from jinjatest.markers import transform_markers
 from jinjatest.rendered import RenderedPrompt
@@ -234,7 +233,8 @@ def from_string(
         if env is None:
             env = create_environment(**env_kwargs)
 
-        instrumentation = instrument(env, test_mode=test_mode)
+        instrumentation = create_instrumentation(test_mode=test_mode)
+        env.globals["jt"] = instrumentation
 
         template = env.from_string(source)
         return cls(
@@ -292,7 +292,8 @@ def from_file(
             template_paths = [template_dir] + [Path(p) for p in template_paths]
 
             env = create_environment(template_paths=template_paths, **env_kwargs)
-            instrumentation = instrument(env, test_mode=test_mode)
+            instrumentation = create_instrumentation(test_mode=test_mode)
+            env.globals["jt"] = instrumentation
         else:
             # For provided env, check if already instrumented
             existing_jt = env.globals.get("jt")
@@ -301,7 +302,8 @@ def from_file(
             ):
                 instrumentation = existing_jt
             else:
-                instrumentation = instrument(env, test_mode=test_mode)
+                instrumentation = create_instrumentation(test_mode=test_mode)
+                env.globals["jt"] = instrumentation
 
         # Determine template name based on how env was obtained
         # When env is provided or template_dir is explicitly set, use full path