turtle_to_stdout.py

#!/usr/bin/env python3
"""Convert an RDF/Turtle CST back to source files, output as JSON array."""

import json
import sys
import xml.etree.ElementTree as ET
from pathlib import Path

from rdflib import Graph, Namespace, RDF
from rdflib.collection import Collection
from pylibsrcml import srcMLArchiveRead


SRCML_NS = "http://www.srcML.org/srcML/src"
CST = Namespace("http://trustgraph.ai/cst#")


def string_to_source_literal(s: str, is_docstring: bool = False) -> str:
    """Convert actual string value back to source representation.

    F-strings are returned unchanged (they were never normalized).
    Docstrings use triple double quotes.
    Regular strings use double quotes to match Python conventions.
    """
    # F-strings were not normalized, return as-is
    if s.startswith(('f"', "f'", 'F"', "F'", 'rf"', "rf'", 'fr"', "fr'")):
        return s

    # Docstrings use triple quotes
    if is_docstring:
        return f'"""{s}"""'

    # Use double quotes, escaping internal double quotes and backslashes
    escaped = s.replace('\\', '\\\\').replace('"', '\\"').replace('\n', '\\n').replace('\r', '\\r').replace('\t', '\\t')
    return f'"{escaped}"'


def dict_to_xml_element(d: dict, ns: str = SRCML_NS) -> ET.Element | str:
    """Convert a dictionary back to an XML element."""
    if d.get("tag") == "#text":
        return d.get("text", "")

    tag = d["tag"]
    full_tag = f"{{{ns}}}{tag}"

    elem = ET.Element(full_tag)

    if "attrs" in d:
        for k, v in d["attrs"].items():
            elem.set(k, v)

    if "text" in d:
        elem.text = d["text"]

    if "children" in d:
        last_elem = None
        for child_dict in d["children"]:
            child = dict_to_xml_element(child_dict, ns)
            if isinstance(child, str):
                if last_elem is not None:
                    last_elem.tail = (last_elem.tail or "") + child
                else:
                    elem.text = (elem.text or "") + child
            else:
                elem.append(child)
                last_elem = child

    return elem


def cst_to_source(cst: dict, filename: str) -> str:
    """Convert a JSON CST back to source code via srcML."""
    unit_elem = dict_to_xml_element(cst)

    unit_elem.set("revision", "1.0.0")
    unit_elem.set("language", "Python")
    unit_elem.set("filename", filename)

    ET.register_namespace("", SRCML_NS)
    root = ET.Element(f"{{{SRCML_NS}}}unit")
    root.set("revision", "1.0.0")
    root.append(unit_elem)

    # Use UTF-8 encoding for proper Unicode support
    xml_bytes = ET.tostring(root, encoding="utf-8", xml_declaration=True)
    xml_str = xml_bytes.decode("utf-8")

    archive = srcMLArchiveRead(xml_str)
    archive.set_src_encoding("UTF-8")
    unit = archive.read_unit()
    if unit is None:
        raise ValueError(f"Failed to read unit for {filename}")

    unit.set_src_encoding("UTF-8")
    source = unit.unparse_string()
    archive.close()

    return source


def rdf_node_to_dict(graph: Graph, node) -> dict:
    """Convert an RDF node back to a CST dict."""
    # Get the type
    node_type = graph.value(node, RDF.type)
    if node_type is None:
        raise ValueError(f"Node {node} has no type")

    type_str = str(node_type)
    if type_str.startswith(str(CST)):
        tag = type_str[len(str(CST)):]
    else:
        tag = type_str

    # Handle TextNode specially
    if tag == "TextNode":
        text = graph.value(node, CST.text)
        return {"tag": "#text", "text": str(text) if text else ""}

    result = {"tag": tag}

    # Get attributes first (need to know if it's a string literal)
    attrs = {}
    for pred, obj in graph.predicate_objects(node):
        pred_str = str(pred)
        if pred_str.startswith(str(CST) + "attr_"):
            attr_name = pred_str[len(str(CST) + "attr_"):]
            attrs[attr_name] = str(obj)

    if attrs:
        result["attrs"] = attrs

    # Get text content, converting string literals back to source form
    text = graph.value(node, CST.text)
    if text is not None:
        text_str = str(text)
        # Convert normalized string back to source literal
        if tag == "literal" and attrs.get("type") == "string":
            is_docstring = attrs.get("format") == "docstring"
            text_str = string_to_source_literal(text_str, is_docstring=is_docstring)
        result["text"] = text_str

    # Get children (rdf:List)
    children_list = graph.value(node, CST.children)
    if children_list is not None:
        children = []
        collection = Collection(graph, children_list)
        for child_node in collection:
            children.append(rdf_node_to_dict(graph, child_node))
        result["children"] = children

    return result


def turtle_to_files(turtle_input: str | Path) -> list[dict]:
    """
    Read Turtle CST and convert back to source files.

    Returns list of {"path": ..., "source": ...} dicts.
    """
    graph = Graph()

    # Load Turtle
    if isinstance(turtle_input, Path) or (
        isinstance(turtle_input, str) and not turtle_input.strip().startswith("@")
    ):
        path = Path(turtle_input)
        if not path.exists():
            print(f"Error: {path} not found", file=sys.stderr)
            sys.exit(1)
        graph.parse(path, format="turtle")
    else:
        graph.parse(data=turtle_input, format="turtle")

    results = []

    # Find all File nodes
    for file_node in graph.subjects(RDF.type, CST.File):
        filepath = str(graph.value(file_node, CST.path))
        root_node = graph.value(file_node, CST.root)

        if root_node is None:
            print(f"Warning: File {filepath} has no root node", file=sys.stderr)
            continue

        cst = rdf_node_to_dict(graph, root_node)
        source = cst_to_source(cst, filepath)

        results.append({
            "path": filepath,
            "source": source
        })

    # Sort by path for consistent output
    results.sort(key=lambda x: x["path"])

    return results


def main():
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <archive.ttl>", file=sys.stderr)
        print(f"  Reads Turtle CST and outputs source files as JSON array", file=sys.stderr)
        sys.exit(1)

    turtle_file = sys.argv[1]
    files = turtle_to_files(turtle_file)

    print(json.dumps(files, indent=2))


if __name__ == "__main__":
    main()