Add an option to flatten a D2 file that contains various imports into a single .d2 file

[behrangsa]

Let's assume that I have:

index.d2 -> imports x.d2, classes.d2
x.d2 -> imports y.d2, classes.d2
y.d2 -> imports z.d2, classes.d2

Add a --flatten <file-name> option to the D2 CLI that flattens index.d2 into a single self-contained file, with all imports reolved, and duplicates (classes.d2) unduplicated.

[baradhili]

I really feel this is not something that should be in the core language. Maybe a bash script/bat file instead?
Though to be honest - imports are not a bad thing and I'm not sure why you would want to flatten everything when you could just zip them?
You would have to recurse through all the import files as well potentially..

but here is a bash script for you to start with

#!/usr/bin/env bash
set -euo pipefail

if [ "$#" -lt 1 ]; then
  echo "Usage: $0 <d2-file> [import-map]" >&2
  exit 1
fi

primary_file="$1"
import_map="${2:-}"

if [ ! -f "$primary_file" ]; then
  echo "Error: $primary_file not found" >&2
  exit 1
fi

dir_of_primary=$(cd "$(dirname "$primary_file")" && pwd)

declare -A import_paths

if [ -s "$import_map" ]; then
  while IFS='=' read -r key value; do
    import_paths["$key"]="$value"
  done < "$import_map"
fi

temp_file=$(mktemp)
trap 'rm -f "$temp_file"' EXIT

while IFS= read -r line; do
  trimmed=${line%%#*}
  trimmed=${trimmed%%$'\r'}
  trimmed=${trimmed//[$'\t\r\n']}

  if [[ $trimmed =~ ^@([[:alnum:]_./-]+)$ ]]; then
    import_key=${BASH_REMATCH[1]}
    resolved=${import_paths[$import_key]:-}

    if [ -z "$resolved" ]; then
      resolved="$dir_of_primary/$import_key"
    fi

    resolved=${resolved%.d2}
    resolved="$resolved.d2"

    if [ -f "$resolved" ]; then
      cat "$resolved"
    else
      echo "Warning: $resolved not found, leaving import line as-is" >&2
      printf '%s\n' "$line"
    fi
  elif [[ $trimmed =~ ^[[:space:]]*([[:alnum:]_./-]+)[[:space:]]*:[[:space:]]*\{[[:space:]]*\.\.\.@[[:alnum:]_./-]+[[:space:]]*\}[[:space:]]*$ ]]; then
    # Spread import inside a map; leave as-is.
    printf '%s\n' "$line"
  else
    printf '%s\n' "$line"
  fi
done < "$primary_file" > "$temp_file"

mv "$temp_file" "$primary_file"

[behrangsa]

I already have written a Python script to do it, but to do it properly I have to write a D2 parser which is not a simple task. Otherwise I have to add hacks to my script to handle edge cases and bugs as they appear.

For example:

• What if I have a single classes file that contains the styling for my graph and then import it into multiple D2 files (for example, in order to individual files display properly when I view them in a D2 viewer inside VS Code)
• What if some of my imports have the .d2 extension and some don't?

Using regular expressions is not the proper way to implement this.

My particular use-case demands a flattened file, otherwise there's nothing wrong with having a main D2 file.

Here's my script:

#!/usr/bin/env python3
"""
D2 Flattener - Expands all imports in a D2 file into a single self-contained file.

Handles:
- Spread imports: ...@path (inlines entire file content)
- Regular imports: @path.node (inlines specific node from file)

Usage:
    python d2-flatten.py input.d2 [-o output.d2]
    python d2-flatten.py input.d2 > output.d2
"""

import argparse
import re
import sys
from pathlib import Path
from typing import Optional


class D2Flattener:
    """Flattens D2 files by expanding all imports inline."""

    # Pattern for spread import: ...@path
    SPREAD_IMPORT_PATTERN = re.compile(r'^(\s*)\.\.\.\s*@([^\s]+)\s*$')

    # Pattern for regular import: identifier: @path.node or identifier: @path
    REGULAR_IMPORT_PATTERN = re.compile(r'^(\s*)(\S+):\s*@([^\s]+)\s*$')

    def __init__(self, base_path: Path):
        self.base_path = base_path
        self.import_chain: set[Path] = set()  # Current import chain for cycle detection
        self.file_cache: dict[Path, str] = {}  # Cache file contents
        self.depth = 0
        self.max_depth = 50  # Prevent infinite recursion

    def flatten(self, file_path: Path) -> str:
        """Flatten a D2 file, expanding all imports."""
        if self.depth > self.max_depth:
            raise RecursionError(f"Max import depth ({self.max_depth}) exceeded")

        resolved_path = self._resolve_path(file_path)

        if not resolved_path.exists():
            raise FileNotFoundError(f"D2 file not found: {resolved_path}")

        # Check for circular imports in current chain only
        if resolved_path in self.import_chain:
            return f"# [Circular import skipped: {file_path}]\n"

        # Return cached content if available (allows multiple imports from same file)
        if resolved_path in self.file_cache:
            return self.file_cache[resolved_path]

        self.import_chain.add(resolved_path)
        self.depth += 1

        try:
            content = resolved_path.read_text()
            result = self._process_content(content, resolved_path.parent)
            self.file_cache[resolved_path] = result  # Cache the result
            return result
        finally:
            self.depth -= 1
            self.import_chain.discard(resolved_path)  # Remove from chain when done

    def _resolve_path(self, file_path: Path) -> Path:
        """Resolve a path, adding .d2 extension if needed."""
        if file_path.is_absolute():
            resolved = file_path
        else:
            resolved = self.base_path / file_path

        # Add .d2 extension if not present
        if resolved.suffix != '.d2':
            resolved = resolved.with_suffix('.d2')

        return resolved.resolve()

    def _process_content(self, content: str, current_dir: Path) -> str:
        """Process file content, expanding imports."""
        lines = content.split('\n')
        result_lines = []

        for line in lines:
            # Check for spread import: ...@path
            spread_match = self.SPREAD_IMPORT_PATTERN.match(line)
            if spread_match:
                indent = spread_match.group(1)
                import_path = spread_match.group(2)
                expanded = self._expand_spread_import(import_path, current_dir, indent)
                result_lines.append(expanded)
                continue

            # Check for regular import: name: @path.node
            regular_match = self.REGULAR_IMPORT_PATTERN.match(line)
            if regular_match:
                indent = regular_match.group(1)
                identifier = regular_match.group(2)
                import_ref = regular_match.group(3)
                expanded = self._expand_regular_import(
                    identifier, import_ref, current_dir, indent
                )
                result_lines.append(expanded)
                continue

            # Not an import, keep line as-is
            result_lines.append(line)

        return '\n'.join(result_lines)

    def _expand_spread_import(
        self, import_path: str, current_dir: Path, indent: str
    ) -> str:
        """Expand a spread import (...@path) by inlining file content."""
        file_path = current_dir / import_path

        # Save and update base path for nested imports
        old_base = self.base_path
        self.base_path = current_dir

        try:
            imported_content = self.flatten(file_path)
            # Add indent to each line
            indented_lines = []
            for line in imported_content.split('\n'):
                if line.strip():  # Only indent non-empty lines
                    indented_lines.append(indent + line)
                else:
                    indented_lines.append(line)
            return '\n'.join(indented_lines)
        except FileNotFoundError as e:
            return f"{indent}# [Import error: {e}]"
        finally:
            self.base_path = old_base

    def _expand_regular_import(
        self, identifier: str, import_ref: str, current_dir: Path, indent: str
    ) -> str:
        """
        Expand a regular import (name: @path.node).

        The import_ref can be:
        - path/to/file.node.subnode (file path + node path)
        - path/to/file (just file path, import entire file as value)
        """
        # Split into file path and node path
        # D2 imports: @dir/file.node means dir/file.d2, node "node"
        parts = import_ref.split('.')
        file_parts = []
        node_parts = []

        # Find where file path ends and node path begins
        # Heuristic: file path contains '/' or is the first part
        for i, part in enumerate(parts):
            if '/' in part or i == 0:
                file_parts.append(part)
            else:
                node_parts = parts[i:]
                break

        file_path_str = '.'.join(file_parts)
        file_path = current_dir / file_path_str

        # Save and update base path for nested imports
        old_base = self.base_path
        self.base_path = current_dir

        try:
            imported_content = self.flatten(file_path)

            if node_parts:
                # Extract specific node from imported content
                node_path = '.'.join(node_parts)
                extracted = self._extract_node(imported_content, node_path)
                if extracted is not None:
                    # Format as: identifier: { extracted_content }
                    return self._format_node_assignment(
                        identifier, extracted, indent
                    )
                else:
                    return f"{indent}# [Node not found: {node_path} in {file_path}]"
            else:
                # Import entire file content
                return self._format_node_assignment(
                    identifier, imported_content, indent
                )
        except FileNotFoundError as e:
            return f"{indent}# [Import error: {e}]"
        finally:
            self.base_path = old_base

    def _extract_node(self, content: str, node_path: str) -> Optional[str]:
        """
        Extract a specific node from D2 content.

        For node_path "sections", finds the "sections: { ... }" block.
        """
        nodes = node_path.split('.')
        current_content = content

        for node_name in nodes:
            extracted = self._extract_single_node(current_content, node_name)
            if extracted is None:
                return None
            current_content = extracted

        return current_content

    def _extract_single_node(self, content: str, node_name: str) -> Optional[str]:
        """Extract a single node's content from D2 text."""
        # Pattern to find node definition: nodename: { or nodename: value
        # This is a simplified parser - handles common cases

        lines = content.split('\n')
        result_lines = []
        brace_depth = 0
        in_target_node = False
        found_node = False
        in_block_string = False  # Track if inside |~...~| block string

        # Pattern for node start: "nodename:" or "nodename: {"
        node_start_pattern = re.compile(
            rf'^(\s*){re.escape(node_name)}\s*:\s*(.*)$'
        )

        for line in lines:
            # Track block string state (|~lang ... ~|)
            # Block strings can span multiple lines
            if not in_block_string and '|~' in line:
                in_block_string = True
                # Check if block string ends on same line
                if '~|' in line and line.index('~|') > line.index('|~'):
                    in_block_string = False
            elif in_block_string and '~|' in line:
                in_block_string = False

            if not in_target_node:
                match = node_start_pattern.match(line)
                if match:
                    found_node = True
                    in_target_node = True
                    remainder = match.group(2).strip()

                    # Check if it's a simple value or block
                    if remainder.startswith('{'):
                        brace_depth = remainder.count('{') - remainder.count('}')
                        # Get content after opening brace
                        after_brace = remainder[1:].strip()
                        if after_brace and not after_brace.startswith('}'):
                            result_lines.append(after_brace)
                        if brace_depth == 0:
                            # Single-line block: node: { content }
                            in_target_node = False
                    elif remainder:
                        # Simple value assignment
                        return remainder
                    else:
                        # node: followed by block on next line
                        brace_depth = 0
            else:
                # Inside target node - only count braces outside block strings
                if in_block_string:
                    # Inside block string, don't count braces
                    open_braces = 0
                    close_braces = 0
                else:
                    open_braces = line.count('{')
                    close_braces = line.count('}')

                if brace_depth == 0 and '{' in line and not in_block_string:
                    # First line with opening brace
                    brace_depth = open_braces - close_braces
                    idx = line.index('{')
                    after_brace = line[idx + 1:].strip()
                    if after_brace and not after_brace.startswith('}'):
                        result_lines.append(after_brace)
                else:
                    brace_depth += open_braces - close_braces

                    if brace_depth > 0:
                        result_lines.append(line)
                    elif brace_depth == 0 and not in_block_string:
                        # Check for content before closing brace
                        if '}' in line:
                            idx = line.index('}')
                            before_brace = line[:idx].strip()
                            if before_brace:
                                result_lines.append(before_brace)
                        in_target_node = False
                        break

        if found_node and result_lines:
            return '\n'.join(result_lines)
        elif found_node:
            return ''  # Empty node
        return None

    def _format_node_assignment(
        self, identifier: str, content: str, indent: str
    ) -> str:
        """Format content as a node assignment."""
        content = content.strip()
        if not content:
            return f"{indent}{identifier}: {{}}"

        # Check if content is already a block or needs wrapping
        lines = content.split('\n')
        if len(lines) == 1 and not content.startswith('{'):
            # Single line, might be simple value
            if ':' in content or '{' in content:
                # Looks like D2 content, wrap in braces
                return f"{indent}{identifier}: {{\n{indent}  {content}\n{indent}}}"
            else:
                return f"{indent}{identifier}: {content}"

        # Multi-line content, wrap in braces with proper indentation
        indented_content = '\n'.join(
            f"{indent}  {line}" if line.strip() else line
            for line in lines
        )
        return f"{indent}{identifier}: {{\n{indented_content}\n{indent}}}"


def main():
    parser = argparse.ArgumentParser(
        description='Flatten D2 files by expanding imports inline.'
    )
    parser.add_argument(
        'input',
        type=Path,
        help='Input D2 file to flatten'
    )
    parser.add_argument(
        '-o', '--output',
        type=Path,
        default=None,
        help='Output file (default: stdout)'
    )

    args = parser.parse_args()

    input_path = args.input.resolve()
    if not input_path.exists():
        print(f"Error: Input file not found: {input_path}", file=sys.stderr)
        sys.exit(1)

    flattener = D2Flattener(input_path.parent)

    try:
        result = flattener.flatten(input_path)

        if args.output:
            args.output.write_text(result)
            print(f"Flattened output written to: {args.output}", file=sys.stderr)
        else:
            print(result)

    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()