API¶

def analyze(source: str, *, module: str = "", file: str = "<source>") -> dict[str, Any]:
    """Run every observation and inference stage over one module.

    Parameters
    ----------
    source : str
        The module's text.
    module : str
        Its dotted import path.
    file : str, optional
        A label for spans.

    Returns
    -------
    dict
        ``facts``, ``names``, ``writes``, ``flow`` and ``plan``: what was seen,
        what it refers to, what was assigned to which typed member, where each
        value came from, and how control moves between steps.
    """
    observed = facts.extract(source, module=module, file=file)
    return {
        "file": file,
        "module": module,
        "facts": observed,
        "names": resolve(observed),
        "writes": resolve_writes(observed),
        "flow": dataflow.analyze(source, module=module, file=file),
        "plan": controlflow.analyze(source, module=module, file=file),
    }

def resolve_module(module: str, origin: str) -> str:
    """Return the absolute path an import names.

    ``from .params import ChargeParam`` inside ``battery.procedure`` names
    ``battery.params``. Extraction records the origin exactly as written,
    because that is what the source says; turning it into a path is an
    inference and belongs here.
    """
    if not origin.startswith("."):
        return origin
    depth = len(origin) - len(origin.lstrip("."))
    parts = module.split(".")[: -depth or None]
    tail = origin.lstrip(".")
    return ".".join([*parts, tail]) if tail else ".".join(parts)

def to_ast_doc(
    source: str,
    *,
    module: str = "",
    profile: str = "ast",
    file: str = "<source>",
    index: dict[str, str] | None = None,
) -> Any:
    """Parse, elide by profile, and collapse resolved constructors.

    Parameters
    ----------
    source : str
        The module's text.
    module : str
        Its dotted import path.
    profile : str, optional
        One of ``awl.vocab.PROFILES``.
    file : str, optional
        A label for spans.
    index : dict, optional
        Module path to source, for the modules this one imports from. A
        parameter object is nearly always defined in another file, so without
        it the collapse almost never fires.

    Returns
    -------
    dict
        An ``AstDoc``. Only constructors whose callee resolved are collapsed;
        an unresolved name stays a plain call.
    """
    observed = facts.extract(source, module=module, file=file)
    return _tree(source, observed, module=module, profile=profile, index=index or {})

def to_compact(
    source: str,
    *,
    module: str = "",
    profile: str = "ast",
    spans: bool = False,
    file: str = "<source>",
    index: dict[str, str] | None = None,
) -> Any:
    """Run the chain and return the editor model.

    Parameters
    ----------
    spans : bool, optional
        Keep source spans, which are the join key for in-place patching and
        for the trace overlay. Not needed to regenerate code.
    """
    observed = facts.extract(source, module=module, file=file)
    tree = _tree(source, observed, module=module, profile=profile, index=index or {}, spans=spans)
    return compact.encode(tree, keep_spans=spans)

def to_document(
    source: str,
    *,
    module: str = "",
    profile: str = "ast",
    file: str = "<source>",
    index: dict[str, str] | None = None,
    layers: tuple[str, ...] | None = None,
    spans: bool | None = None,
    identities: bool | None = None,
    trivia: bool | None = None,
) -> dict[str, Any]:
    """Build the JSON-LD document a profile calls for.

    Parameters
    ----------
    source : str
        The module's text.
    module : str
        Its dotted import path.
    profile : str, optional
        A named set of generator parameters: which wrappers are transparent,
        which types go opaque, whether keywords fold, and which lookups run.
        Defaults to ``ast``, the profile whose obligation is complete value
        provenance. The reduced profiles are paused, and until each is defined
        by the class of question it must answer, choosing one only makes the
        document smaller and no better.
    file : str, optional
        A label for spans.
    index : dict, optional
        Module path to source, for the modules this one imports from.
    layers : tuple of str, optional
        Overrides the profile's lookups, for a caller who wants one of them on
        its own. Asking for ``("document",)`` gives the tree and nothing
        derived from it, which is what a reader comparing notations needs.
    spans : bool, optional
        Overrides the profile's :data:`awl.vocab.MATERIALIZES_SPANS`. On for
        every profile, because a span locates a node in the file it came from,
        which is what a patch and a trace overlay are written against.
    identities : bool, optional
        Overrides the profile's :data:`awl.vocab.MATERIALIZES_IDENTITIES`. On
        for every profile: it names each statement with the identity the plan
        mints for it, so the tree's statement and the plan's step are one node
        rather than two that happen to sit at the same coordinates. Needs
        ``spans``, which is what matches the two sides at build time.
    trivia : bool, optional
        Overrides the profile's :data:`awl.vocab.MATERIALIZES_TRIVIA`. On for
        every profile: it carries the comment written about each statement,
        which the syntax tree has no node for. Needs ``spans`` too, since a
        comment is placed against the span of the statement it describes.

    Returns
    -------
    dict
        A ``@context`` and a ``@graph``. This is the artefact; RDF is one
        serialization of it and the editor model is the first entry in it.

    Raises
    ------
    ValueError
        If a layer is not one of :data:`awl.vocab.LAYERS`, rather than quietly
        returning less than was asked for.
    """
    selected = vocab.LOOKUPS[profile] if layers is None else layers
    located = vocab.MATERIALIZES_SPANS[profile] if spans is None else spans
    named = vocab.MATERIALIZES_IDENTITIES[profile] if identities is None else identities
    noted = vocab.MATERIALIZES_TRIVIA[profile] if trivia is None else trivia
    unknown = set(selected) - set(vocab.LAYERS)
    if unknown:
        raise ValueError(f"unknown layers {sorted(unknown)}; expected some of {list(vocab.LAYERS)}")

    observed = facts.extract(source, module=module, file=file)
    # The same resolution the collapse used. Building the context from this
    # module's types alone left an imported class resolving through @vocab, so
    # the node was typed awl:ChargeParam while the collapse had named it
    # py/tier3_oold.params/ChargeParam: one entity, two IRIs.
    types, _ = _collapsible(observed, module, index or {})
    built = context.build_context(list(types.values()))

    graph: list[Any] = []
    for part in _layers(
        source,
        observed,
        selected,
        module=module,
        profile=profile,
        file=file,
        index=index,
        spans=located,
        # A statement is matched to its step by span, so without one there is
        # nothing to name and asking for both is asking for neither.
        identities=named and located,
        # Same bargain: a comment is placed against the span of the statement
        # it was written about.
        trivia=noted and located,
    ):
        graph.extend(part.get("@graph", [part]))
    return {"@context": built["@context"], "@graph": graph}

def to_graph(
    source: str,
    *,
    module: str = "",
    profile: str = "ast",
    file: str = "<source>",
    index: dict[str, str] | None = None,
    layers: tuple[str, ...] | None = None,
    spans: bool | None = None,
    identities: bool | None = None,
    trivia: bool | None = None,
):
    """Serialize :func:`to_document` as RDF, taking the same parameters.

    Returns
    -------
    rdflib.Graph

    Notes
    -----
    There is nothing here but a change of notation. What the graph contains is
    decided by the profile when the document is built, so the two cannot say
    different things about one program.
    """
    return rdf.to_graph(
        to_document(
            source,
            module=module,
            profile=profile,
            file=file,
            index=index,
            layers=layers,
            spans=spans,
            identities=identities,
            trivia=trivia,
        )
    )

def trace_run(
    source: str,
    call,
    *,
    module: str = "",
    file: str = "<source>",
) -> dict[str, Any]:
    """Run *call* under instrumentation and join the events onto the plan.

    Parameters
    ----------
    source : str
        The text of the module being run, so its plan can be built.
    call : callable
        Invoked with no arguments.

    Returns
    -------
    dict
        One entry per step: whether it ran, on which loop iterations, and for a
        branch which outcomes were observed.
    """
    from awl.trace import trace

    events = trace(call)
    plan = controlflow.analyze(source, module=module, file=file)
    return execution.join(plan, events)

def extract(source: str, *, module: str, file: str = "<source>") -> dict[str, Any]:
    """Read symbol facts and type info out of one module.

    Parameters
    ----------
    source : str
        The module's text.
    module : str
        Its dotted import path, used to mint identities.
    file : str, optional
        A label for spans.

    Returns
    -------
    dict
        Conforms to ``symbol-facts.schema.json``. ``types`` is populated only
        for classes deriving from ``LinkedBaseModel``; a plain dataclass yields
        a declaration with fields and no ``TypeInfo``, which is the tier 2 rung.

    Notes
    -----
    Never imports the module, and never decides what a name refers to.
    """
    walk = _Walk(module, file)
    walk.visit(ast.parse(source))
    return {
        "file": file,
        "module": module,
        "declarations": walk.declarations,
        "imports": walk.imports,
        "aliases": walk.aliases,
        "exports": walk.exports,
        "uses": walk.uses,
        "writes": walk.writes,
        "bindings": walk.bindings,
        "types": walk.types,
    }

def resolve(
    facts: dict[str, Any],
    *,
    index: dict[str, Any] | None = None,
    scheme: str = "py",
) -> dict[str, Any]:
    """Bind each use to an identity with a confidence tier.

    Parameters
    ----------
    facts : dict
        A ``SymbolFacts`` document.
    index : dict, optional
        Module path to that module's ``SymbolFacts``. Supplying it lets a
        re-export be followed one hop, which is a deduction and is therefore
        marked ``INFERRED``. Without it, resolution is per file and an import
        binds to the module named at the import site, which is what the source
        literally says.
    scheme : str, optional
        Language dimension passed through to minting.

    Returns
    -------
    dict
        Conforms to ``resolved-names.schema.json``. A binding with no identity
        is ``AMBIGUOUS``, never a fabricated identity: a wrong identity merges
        two entities, which is worse than admitting ignorance.
    """
    index = index or {}
    imports = facts.get("imports", [])
    by_name = {entry["local_name"]: entry for entry in imports}
    declarations = facts.get("declarations", [])
    declared = {entry.get("name") for entry in declarations}
    has_star = any(entry.get("is_star") for entry in imports)
    module = facts.get("module", "")

    bindings = []
    for use in facts.get("uses", []):
        name = use["local_name"]
        identity: dict[str, Any] | None = None

        if _is_overloaded(declarations, name):
            confidence = AMBIGUOUS
        elif name in by_name:
            identity, confidence = _bind_import(by_name[name], index, scheme)
        elif name in declared:
            identity = mint(scheme=scheme, module=module, symbol=name)
            confidence = EXTRACTED
        else:
            # Reached through a star import, or simply not in the scanned set.
            # Nothing in the surveyed prior art resolves a star import soundly,
            # so it stays flagged and the domain schema forbids it.
            confidence = AMBIGUOUS

        binding: dict[str, Any] = {
            "local_name": name,
            "span": use["span"],
            "confidence": confidence,
        }
        if identity is not None:
            binding["identity"] = identity
        if confidence is AMBIGUOUS and has_star:
            binding["reason"] = "reached through a star import"
        bindings.append(binding)

    return {"file": facts["file"], "bindings": bindings}

def resolve_writes(
    facts: dict[str, Any],
    *,
    scheme: str = "py",
    imported: dict[str, dict[str, Any]] | None = None,
) -> dict[str, Any]:
    """Resolve each attribute write to the member it targets.

    Parameters
    ----------
    facts : dict
        A ``SymbolFacts`` document, whose ``writes`` carry attribute paths,
        whose ``bindings`` carry local aliases, and whose ``declarations``
        carry the annotations to walk them against.
    scheme : str, optional
        Language dimension passed through to minting.
    imported : dict, optional
        Classes declared in the modules this one imports from, by name. A
        write walks declared fields, and a declaration lives in the module that
        made it, so without these an imported class resolves for the collapse
        and not for the write: known and unknown in the same pass.

    Returns
    -------
    dict
        ``{"file": ..., "writes": [...]}``. Each entry carries the path as
        written, the canonical member path, the member assigned, the class
        declaring it, the class the chain started from, the range, and a
        confidence tier.

    Notes
    -----
    This turns "an attribute named ``e_mod`` was assigned at line 74" into "the
    modulus of elasticity of a tensile test specimen, reached from a tensile
    test dataset, was assigned at line 74, by ``ModulusOfElasticity.from_pint``".

    Bindings and writes are replayed in **source order** within each function,
    so a local alias is in scope for the writes that follow it and the rooted
    chain is carried across it. ``s = dataset.specimen`` followed by ``s.e_mod
    = ...`` therefore produces the same member, owner and root as the direct
    form, differing only in confidence.

    A hop that cannot be walked yields ``AMBIGUOUS`` rather than a guess.
    """
    # A local declaration wins: a class declared here is what a name here
    # means, whatever a module it imports from happens to call the same thing.
    classes = {**(imported or {}), **_classes(facts)}
    declared = {
        name: entry["identity"]["iri"] for name, entry in classes.items() if (entry.get("identity") or {}).get("iri")
    }
    imports = {entry["local_name"]: entry for entry in facts.get("imports", [])}
    module = facts.get("module", "")

    def identity_for(class_name: str) -> str:
        """Return the identity of a class, as whoever declared it minted it.

        The declaration is asked first. Re-minting from the import as written
        gets the module wrong for anything reached indirectly: the range of an
        imported class's field is declared where that class is, not where it
        was imported to.
        """
        if class_name in declared:
            return declared[class_name]
        entry = imports.get(class_name)
        if entry is not None:
            return mint(scheme=scheme, module=entry["from_module"], symbol=class_name)["iri"]
        return mint(scheme=scheme, module=module, symbol=class_name)["iri"]

    def position(event: dict[str, Any]) -> tuple[int, int]:
        span = event.get("span") or {}
        return span.get("start_line", 0), span.get("start_col", 0)

    events: list[dict[str, Any]] = [
        *({"kind": "binding", **entry} for entry in facts.get("bindings", [])),
        *({"kind": "write", **entry} for entry in facts.get("writes", [])),
    ]
    events.sort(key=position)

    environments: dict[str | None, dict[str, _Bound]] = {}
    for entry in facts.get("declarations", []):
        if entry.get("kind") != "function":
            continue
        environments[entry["name"]] = {
            parameter["name"]: _Bound(target, EXTRACTED)
            for parameter in entry.get("parameters", [])
            if (target := annotation_target(parameter.get("annotation")))
        }

    resolved = []
    for event in events:
        environment = environments.setdefault(event.get("in_function"), {})

        if event["kind"] == "binding":
            bound = _bound_from(event, environment, classes)
            if bound is None:
                # A rebinding to something unknown drops the old type rather
                # than leaving a stale one in scope.
                environment.pop(event["name"], None)
            else:
                environment[event["name"]] = bound
            continue

        resolved.append(_describe(event, _walk(event["path"], environment, classes), identity_for))

    return {"file": facts["file"], "writes": resolved}

def elide(doc: Any, *, profile: str = "ast", source: str = "") -> Any:
    """Apply a cutoff profile to an AST document.

    Parameters
    ----------
    doc : dict or list
        An ``AstDoc``.
    profile : str
        One of ``awl.vocab.PROFILES``.
    source : str, optional
        Original text, used to populate ``source_text`` on opaque nodes. Without
        it an opaque node keeps its type but loses its expression.

    Returns
    -------
    dict or list
        A new document; the input is not mutated.

    Raises
    ------
    KeyError
        If the profile is unknown, rather than silently doing nothing.
    """
    return _walk(doc, TRANSPARENT[profile], OPAQUE[profile], FOLDS_KEYWORDS[profile], source)

def rewrap_statements(node: dict[str, Any]) -> dict[str, Any]:
    """Put back the ``Expr`` wrappers that elision dropped.

    Python needs a statement to hold an expression in a body. *Which* items
    need one is entirely derivable — anything in a statement list that is not
    itself a statement — so dropping the wrapper costs nothing and restoring
    it needs no record of what was removed. That is what makes ``Expr``
    elidable on the profile that regenerates code, where a lossy elision would
    not be.
    """
    statements = statement_types()
    out = dict(node)
    for field in ORDERED_FIELDS:
        items = out.get(field)
        if not isinstance(items, list):
            continue
        out[field] = [
            item
            if not isinstance(item, dict) or item.get("_type") in statements or "_type" not in item
            else {
                "_type": "Expr",
                "value": item,
                **{key: item[key] for key in _POSITION if key in item},
            }
            for item in items
        ]
    return out

def unfold(doc: Any) -> Any:
    """Restore everything elision rewrote reversibly.

    Parameters
    ----------
    doc : dict or list
        An ``AstDoc`` produced by :func:`elide`.

    Returns
    -------
    dict or list
        A document whose calls carry ``keywords`` again and whose bodies carry
        their statement wrappers, ready to unparse.
    """
    if isinstance(doc, list):
        return [unfold(item) for item in doc]
    if not isinstance(doc, dict):
        return doc
    return rewrap_statements(unfold_node({key: unfold(value) for key, value in doc.items()}))

def unfold_node(node: dict[str, Any], *, type_key: str = "_type") -> dict[str, Any]:
    """Reverse keyword folding for one node.

    Parameters
    ----------
    node : dict
        A node that may carry ``keyword_arguments``.
    type_key : str, optional
        The key naming a node type. The compact form spells it ``type`` and
        the intermediate form ``_type``; parameterising it keeps one
        implementation of how folding reverses, rather than two that drift.

    Returns
    -------
    dict
        The node with its ``keywords`` list restored, or unchanged.

    Notes
    -----
    Folding is a rewrite, not a loss, and this is what makes that true. Every
    profile folds because of it, including the faithful one: without a working
    inverse the profile that regenerates code could not fold, and the
    constructor collapse could never fire where it is most useful.
    """
    folded = node.get("keyword_arguments")
    if not isinstance(folded, dict):
        return node
    out = {key: value for key, value in node.items() if key != "keyword_arguments"}
    position = {key: node[key] for key in _POSITION if key in node}
    # Anything already here had no name to fold under, `**kwargs` being the
    # only case. Overwriting the list rather than extending it dropped it.
    survived = list(out.get("keywords", []) or [])
    out["keywords"] = [
        {
            type_key: "keyword",
            "arg": name,
            "value": {k: v for k, v in value.items() if k not in _ORDERING} if isinstance(value, dict) else value,
            **position,
        }
        for name, value in folded.items()
    ] + survived
    return out

def callee_of(node: dict[str, Any]) -> str | None:
    """Return the local class name a collapsed node names, or None.

    Read from ``@type`` when it carries a bare term. A term resolves through
    the context and is therefore both the class name and, once mapped, the
    IRI; a CURIE or absolute IRI is emitted verbatim and names no local class.
    That distinction is what removes the need for a separate ``_callee``.
    """
    declared = node.get("@type")
    if isinstance(declared, str):
        declared = [declared]
    for candidate in declared or []:
        if isinstance(candidate, str) and ":" not in candidate and "/" not in candidate:
            return candidate
    return None

def collapse(
    doc: Any,
    *,
    types: dict[str, Any],
    resolved: dict[str, str],
    embed_context: bool = True,
    keep_spans: bool = False,
) -> Any:
    """Rewrite resolved constructor calls into typed nodes.

    Parameters
    ----------
    doc : dict or list
        An ``AstDoc``.
    types : dict
        ``TypeInfo`` keyed by symbol name.
    resolved : dict
        Local name to symbol name. A name absent here is not collapsed, which
        is how an ambiguous binding stays a plain call: a wrong type is worse
        than no type.
    embed_context : bool, optional
        Emit a per-node ``@context``. Turn it off when the document already
        carries a context naming these types, which is the compact form: a
        node is then just its class name and its data.
    keep_spans : bool, optional
        Emit the source span as ``"@"``. Needed only to patch the original
        file in place; regenerating code from the document does not use it.
        Off by default, matching the compact encoder.

    Returns
    -------
    dict or list
        A new document.
    """
    options: dict[str, Any] = {
        "types": types,
        "resolved": resolved,
        "embed_context": embed_context,
        "keep_spans": keep_spans,
    }
    if isinstance(doc, list):
        return [collapse(item, **options) for item in doc]
    if not isinstance(doc, dict):
        return doc

    if doc.get("_type") == "Call":
        callee = doc.get("func", {}).get("id")
        symbol = resolved.get(callee) if callee else None
        info = types.get(symbol) if symbol else None
        if info is not None:
            return _typed_node(doc, info, callee, **options)

    return {key: collapse(value, **options) for key, value in doc.items()}

def expand(node: Any) -> Any:
    """Rebuild a ``Call`` document from a typed node.

    Parameters
    ----------
    node : dict or list
        A document that may contain collapsed nodes.

    Returns
    -------
    dict or list
        An ``AstDoc`` with every typed node turned back into a call.

    Raises
    ------
    ValueError
        If a typed node carries no ``_callee``. Emitting a call with a guessed
        name would be worse than refusing.

    Notes
    -----
    Field insertion order is the keyword order, which JSON objects and Python
    dictionaries both preserve, so the restored call reads as it was written.
    """
    if isinstance(node, list):
        return [expand(item) for item in node]
    if not isinstance(node, dict):
        return node
    callee = callee_of(node)
    if "@type" in node and callee is None:
        raise ValueError("typed node names no local class, so it cannot be expanded; expected a bare term in @type")
    if callee is None:
        expanded = {key: expand(value) for key, value in node.items()}
        return unfold_node(expanded)

    span = node.get("span") or [1, 0, 1, 0]
    position = {
        "lineno": span[0],
        "col_offset": span[1],
        "end_lineno": span[2],
        "end_col_offset": span[3],
    }
    keywords = []
    for name, value in node.items():
        if name in _RESERVED:
            continue
        inner = expand(value) if isinstance(value, dict) else {"_type": "Constant", "value": value, **position}
        keywords.append({"_type": "keyword", "arg": name, "value": inner, **position})

    return {
        "_type": "Call",
        "func": {"_type": "Name", "id": callee, "ctx": {"_type": "Load"}, **position},
        "args": [],
        "keywords": keywords,
        **position,
    }

def decode(node: Any) -> Any:
    """Decode the compact form back into a live AST node.

    Parameters
    ----------
    node : dict or list
        A ``CompactDoc``.

    Returns
    -------
    ast.AST or list
        Ready for :func:`ast.fix_missing_locations` and :func:`ast.unparse`.

    Notes
    -----
    Every field of the target class is rebuilt from ``cls._fields``, which is
    what makes the encoder's omissions safe. Copying only the present keys
    works on a recent interpreter and breaks on an older one: 3.11 raises
    ``AttributeError`` unparsing a ``Module`` with no ``type_ignores`` where
    3.13 returns the source.
    """
    if isinstance(node, list):
        return [decode(item) for item in node]
    if not isinstance(node, dict):
        return node

    if "@type" in node and not _is_ast_node(node):
        return _decode_collapsed(node)
    # Minus the drop list, because a shorthand can arrive carrying a
    # materialized ordering: renumbering a statement list stamps `order` onto
    # every item, and a docstring is an item. Reading that as a typed node
    # asked a literal for its `@type` and raised, which made adding a step to
    # any body holding a docstring fail on the way back out.
    plain = set(node) - _DROP
    if "literal" in node and plain <= {"literal", *_SHORTHAND_EXTRA}:
        from awl.astdoc import from_doc

        return ast.Constant(value=from_doc(node["literal"]))
    if "var" in node and plain <= {"var", *_SHORTHAND_EXTRA}:
        return ast.Name(id=node["var"], ctx=ast.Load())

    from awl.elide import unfold_node

    node = unfold_node(node, type_key="@type")
    cls = getattr(ast, node["@type"])
    return cls(**{field: _decode_field(node, cls, field) for field in cls._fields})

def dumps(doc: Any, *, width: int = 88, indent: int = 1, _depth: int = 0) -> str:
    """Serialize a compact document, inlining whatever fits.

    Parameters
    ----------
    doc : Any
        A ``CompactDoc``, or any JSON-serializable value.
    width : int, optional
        The column a line may reach before its structure is broken open.
    indent : int, optional
        Spaces per level.

    Returns
    -------
    str
        JSON in which a small structure stays on one line and a large one
        breaks.

    Notes
    -----
    ``json.dumps(indent=...)`` puts every element of every structure on its own
    line, so ``{"var": "i"}`` costs three lines and a two-argument call costs a
    page. That is not more readable, only taller: the shape of a node is
    easiest to see when the node fits on one line.
    """
    pad = " " * (indent * _depth)
    inner = " " * (indent * (_depth + 1))

    if isinstance(doc, dict):
        if not doc:
            return "{}"
        flat = (
            "{"
            + ", ".join(json.dumps(key) + ": " + dumps(value, width=width, indent=0) for key, value in doc.items())
            + "}"
        )
        if len(pad) + len(flat) <= width and NEWLINE not in flat:
            return flat
        parts = [
            inner + json.dumps(key) + ": " + dumps(value, width=width, indent=indent, _depth=_depth + 1)
            for key, value in doc.items()
        ]
        return "{" + NEWLINE + ("," + NEWLINE).join(parts) + NEWLINE + pad + "}"

    if isinstance(doc, list):
        if not doc:
            return "[]"
        flat = "[" + ", ".join(dumps(item, width=width, indent=0) for item in doc) + "]"
        if len(pad) + len(flat) <= width and NEWLINE not in flat:
            return flat
        parts = [inner + dumps(item, width=width, indent=indent, _depth=_depth + 1) for item in doc]
        return "[" + NEWLINE + ("," + NEWLINE).join(parts) + NEWLINE + pad + "]"

    return json.dumps(doc)

def encode(doc: Any, *, keep_spans: bool = False) -> Any:
    """Encode an ``AstDoc`` into the compact form.

    Parameters
    ----------
    doc : dict or list
        An ``AstDoc``, from ``ast2json`` or from the elision stage.
    keep_spans : bool, optional
        Keep a ``span`` array. The editor addresses edits by span
        and the trace overlay attributes events by span, so without it there is
        no join key between the editor, the trace and the source. Costs about
        two percentage points of size, so the RDF projection leaves it off.

    Returns
    -------
    dict or list
        A ``CompactDoc``.

    Notes
    -----
    Rules: ``_type`` becomes ``type``; a bare ``Constant`` becomes
    ``{"literal": value}``; a bare ``Name`` becomes ``{"var": id}``; an
    operator becomes its bare name; ``ctx``, positions, nulls, empty lists and
    the derivable orderings are dropped.
    """
    if isinstance(doc, list):
        return [encode(item, keep_spans=keep_spans) for item in doc]
    if not isinstance(doc, dict):
        return doc

    node = _shorthand(doc)
    if node is None:
        return _encode_long(doc, keep_spans=keep_spans)

    if keep_spans and "lineno" in doc:
        node["span"] = _span(doc)
    return node