60 lines
1.5 KiB
Python
60 lines
1.5 KiB
Python
from typing import Any, Collection, Dict, List, Optional, overload
|
|
|
|
from ..language import Node, OperationType
|
|
from ..pyutils import is_iterable
|
|
|
|
__all__ = ["ast_to_dict"]
|
|
|
|
|
|
@overload
|
|
def ast_to_dict(
|
|
node: Node, locations: bool = False, cache: Optional[Dict[Node, Any]] = None
|
|
) -> Dict: ...
|
|
|
|
|
|
@overload
|
|
def ast_to_dict(
|
|
node: Collection[Node],
|
|
locations: bool = False,
|
|
cache: Optional[Dict[Node, Any]] = None,
|
|
) -> List[Node]: ...
|
|
|
|
|
|
@overload
|
|
def ast_to_dict(
|
|
node: OperationType,
|
|
locations: bool = False,
|
|
cache: Optional[Dict[Node, Any]] = None,
|
|
) -> str: ...
|
|
|
|
|
|
def ast_to_dict(
|
|
node: Any, locations: bool = False, cache: Optional[Dict[Node, Any]] = None
|
|
) -> Any:
|
|
"""Convert a language AST to a nested Python dictionary.
|
|
|
|
Set `locations` to True in order to get the locations as well.
|
|
"""
|
|
if isinstance(node, Node):
|
|
if cache is None:
|
|
cache = {}
|
|
elif node in cache:
|
|
return cache[node]
|
|
cache[node] = res = {}
|
|
res.update(
|
|
{
|
|
key: ast_to_dict(getattr(node, key), locations, cache)
|
|
for key in ("kind",) + node.keys[1:]
|
|
}
|
|
)
|
|
if locations:
|
|
loc = node.loc
|
|
if loc:
|
|
res["loc"] = dict(start=loc.start, end=loc.end)
|
|
return res
|
|
if is_iterable(node):
|
|
return [ast_to_dict(sub_node, locations, cache) for sub_node in node]
|
|
if isinstance(node, OperationType):
|
|
return node.value
|
|
return node
|