Package: pyoak
Commerial-grade, well tested, documented and typed Python library for modeling, building, traversing, transforming, transferring and even pattern matching abstract syntax trees (AST) for arbtirary languages.Features🌳 Easy to use dataclasses based, strictly typed, pseudo-immutable AST modeling📝 "Magic", auto-maintained node registry, allowing node cross-referncing and retrievel📚 Source origin tracking for each node📺 Zero-setup pretty printing usingRich💾 json, msgpack, yaml or plain dict (de)serialization🏃‍♀️ AST traversal: depth-first or breadth-first, top down or bottom up, with filtering and pruning🎯 Xpath-like AST search (top to the node)👯‍♂️ Node pattern matching with ability to capture specific sub-trees or attributes... and more!Feature RoadmapMake orjson, pyyaml, msgpack & chardet optional dependenciesrustify some parts for performancewell, that's too far-fetchedTable of ContentsInstallationBasic UsageDocumentationDefining a ModelInheritanceCreating a NodeRuntime Type ChecksID's & RegistryMutatingApplying Changes to a NodeTraversalGoing Down the TreeGoing Up the TreeOther HelpersXpath & Pattern MatchingXpathPattern MatchingVisitors & TransformersASTVisitorASTTransformVisitorSerialization & DeserializationDebugging & Reporting Deserialization Errors for Deeply Nested TreesCreditsLinksInstallationpipinstallpyoakBasic Usagefromdataclassesimportdataclassfrompyoak.nodeimportASTNodefrompyoak.originimportCodeOrigin,TextSource,get_code_rangefromrichimportprint@dataclass(frozen=True)className(ASTNode):identifier:str@dataclass(frozen=True)classNumberLiteral(ASTNode):value:int@dataclass(frozen=True)classStmt(ASTNode):pass@dataclass(frozen=True)classAssignStmt(Stmt):name:Namevalue:NumberLiteraldefparse_assignment(code:str)->AssignStmt:# This should be a real parse logics=TextSource(source_uri="important_place",source_type="text/plain",_raw=code)stmt_o=CodeOrigin(source=s,position=get_code_range(0,1,0,5,1,5))x_o=CodeOrigin(source=s,position=get_code_range(0,1,0,1,1,1))num_o=CodeOrigin(source=s,position=get_code_range(4,1,4,5,1,5))returnAssignStmt(name=Name(identifier="x",origin=x_o),value=NumberLiteral(value=1,origin=num_o),origin=stmt_o,)node=parse_assignment("x = 1\nother_stuff")print(node)# Prints a pretty treeprint(f"Original source code:{node.origin.get_raw()}")# prints `Original source code: x = 1`assertnodeisAssignStmt.get(node.id)# you can always get the node using only its iddup_node=node.duplicate()# returns a deep copy of the nodeassertnodeisnotdup_node# they are different, including id, but...assertdup_node==node# True# which is the same asassert(dup_node.content_id==node.content_id)# content id uniquely represents the subtree# Now let's get an iterable of all Name, NumberLiteral nodes# this will traverse all the way down the tree, not just the direct childrensome_children=node.gather((Name,NumberLiteral))forsubtreeinsome_children:print(subtree)print(f"Original source code:{subtree.origin.get_raw()}")and this is just the tip of the iceberg!DocumentationThe docs are not a book on code parsing, semantic analysis or AST transformations. It is assumed that you already know what ASTs are and are looking for a library to make it a bit easier and more convenient to work with them in Python.pyoak code is strictly typed and heavily documented. Most public (and private) methods and system fields have extensive docstrings, thus it should be easy to use the library without reading the docs.It is strongly encouraged to use mypy/pyright (or your type checker of choice) as it will make working with ASTs a breeze, as well as help you to avoid many common mistakes.Defining a Modelpyoak AST nodes are regular Pythondataclasses, witha lot ofsome additional magic.To create a node you just need to create a frozen dataclass that inherits frompyoak.node.ASTNode.fromdataclassesimportdataclassfrompyoak.nodeimportASTNode@dataclass(frozen=True)classMyNode(ASTNode):attribute:intchild:ASTNode|Nonemore_children:tuple[ASTNode,...]All fields are classified into either child fields or properties. Any subclass of ASTNode in field type marks it as a child field.Since nodes are assumed to be immutable, pyoak will check type annotations on the best effort basis, to prevent usage of mutable types by raising an exception at class defintion. The following rules are enforced:Child fields may only be a union ofASTNodesubclasses (including optionalNone), or a tuple ofASTNodesubclasses. Anything else will raise an exception.For properties, nutable collections are not allowed, but arbitrary types are. It is up to you to ensure that the type is immutable.💡 You may use dataclasses.field, e.g. to define defaults, provide serialization hints via metadata, or mark some fields as non-compare. pyoak doesn't restrict you in any wayInheritanceIt may be convenient to use inheritance. Especially when combined with the fact that dataclasses can override fields of their parent classes. See the example below:@dataclass(frozen=True)classExpr(ASTNode):"""Base for all expressions"""pass@dataclass(frozen=True)classLiteral(Expr):"""Base for all literals"""value:Any@dataclass(frozen=True)classNumberLiteral(Literal):value:int@dataclass(frozen=True)classNoneLiteral(Literal):value:None=field(default=None,init=False)@dataclass(frozen=True)classBinExpr(Expr):"""Base for all binary expressions"""left:Exprright:ExprNot only have you defined the model for binary expressions, but also:narrowed down the value type forNumberLiteraltointsaved yourself time by not having to definevaluefield forNoneLiteralwhen youinstantiateall while maintaining full class compatibility.💡 You may use multiple inheritance as well, say, to define "tag" like classes, that do not carry any data, but are used to mark nodes. E.g.Unsupportedfor structures that wrap unparsed code. However, this will prevent you from using slotted classesCreating a NodeIn its simplest form you just instantiate a class:int_val=NumberLiteral(value=1)This will create a node with automatically generated id.The base ASTNode class defines only one optional fieldorigin(which defaults toNO_ORIGINsentinel)int_val=NumberLiteral(value=1,origin=CodeOrigin(...))Runtime Type Checkspyoak has an optional feature (disabled by default) to do runtime type checks on instance creation. To enable it use the following snippet:frompyoakimportconfigconfig.RUNTIME_TYPE_CHECK=TrueThis feature is convenient during development, to check for typical mistakes, like passing list instead of tuple. However, it is not recommended to use it in production, as it slows down instance creation by 20-40%.ID's & RegistryEvery time you create a node, it gets two ids:idandcontent_id. The former is a globally unique id for the node, the latter is a hash representing node's content (itself and all of it's children).Each new node is automatically added to the registry, which is a mapping fromidto the node. This allows you to retrieve any node by its id, even if you don't have a reference to it.int_val=NumberLiteral(value=1)assertint_valisNumberLiteral.get(int_val.id)# orassertint_valisASTNode.get_any(int_val.id)However, registry is a weak mapping, so if you don't have any references to the node anywhere, it will be garbage collected and removed from the registry.MutatingApplying Changes to a NodeNodes are immutable (to the extent possible in Python), thus if you need to change values of a node, usedataclasses.replacefunction.TraversalGoing Down the TreeEach node, knows about its children:forcinnode.get_child_nodes():print(c)# orlist_of_children=node.childrenThere are multiple helpers to traverse down the tree:fornum_litsinnode.dfs(skip_self=False,filter=lambdan:isinstance(n,NumberLiteral)):print(num_lits.value)# or using a shorthandfornum_litsinnode.gather(NumberLiteral):print(num_lits.value)# breadth-first search also availablelist(node.bfs(filter=important_nodes,prune=stop_at))Note, that these methods return an generator, so you can stop the traversal at any time. Also, worth mentioning that the node itself is not yielded.Going Up the TreeIf you need to search "up" the tree, you need to use theTreeinstance created from the root node:tree=Tree(root_node)# ortree=root_node.to_tree()Now you can use multiple methods:parent=tree.get_parent(node)# to get the parent of the nodeforancestorintree.get_ancestors(node):print(ancestor)parent_stmt=tree.get_first_ancestor_of_type(expr_node,Stmt)Other HelpersBesides the traversal methods mentioned above, there are also some other helpers:ASTNode.duplicate- create a deep copy of the given subtree, recursively duplicating all of its childrenASTNode.to_properties_dict- that gives a mapping of field names to values only for the properties of the node, without the childrenTree.is_ancestor- check if a node is an ancestor of another nodeTree.get_depth- get the depth of the node in the tree, with an optionalrelative_toparameter to get the depth relative to another nodeand others, that are less commonly used.Xpath & Pattern Matchingpyoak.match package provides two ways to match nodes in the tree: "xpath" and pattern matching.Xpath can be used to match nodes by their path from the root of the tree using a syntax similar to xpath in xml.Pattern matching is concerned with matching nodes by their content, including their children.XpathIf you want to find all sub-nodes matching a given xpath starting a given node, you can use thefindandfindallmethods:root_node=parse("some code")fornodeinroot_node.findall("//IntLiteral"):print(node.value)If you want to match a given node, i.e. check if it's in the path, you must create an ASTXpath object and call itsmatchmethod:frompyoak.match.xpathimportASTXpathroot_node=parse("some code")xpath=ASTXpath("//IntLiteral")...ifxpath.match(root_node,some_node):print("Matched!")SyntaxEach path is in the format@parent_field_name[index]type_nameand is separated by a slash.The slash is optional at the start of the XPath. If omitted, it is the same as using//the rest of the path.//is a wildcard (anywhere) that matches any path.@parent_field_nameandindexare optional.type_nameis optional, except for the last path in the XPath.Types are instance comparisons, so any subclass matches a type.More examplesCallExpr/@arguments[0]Idmatches the first argument of a call expression at any level in the tree, which is of typeId.CallExpr/@arguments[*]Idmatches any argument of typeId.@arguments[*]CallExprmatches any CallExpr as long as it is an argument (stored in a field with the said name) of a parent expression.CallExprmatches any call expression at any level in the tree.Pattern MatchingPattern matching is done using a "matcher" object. There are two of them:MultiPatternMatcherfor matching against multiple patternsNodeMatcherfor matching against a single patternTo create aNodeMatcher:macher,msg=NodeMatcher.from_pattern("(RootClass @child_tuple=[(*) -> cap $cap *])")assertmacherisnotNone,msgand then match:node=RootClass(child_tuple=(c1,c2,tail)ok,match_dict=macher.match(node)assertokassertmatch_dict["cap"]==f1WithMultiPatternMatcheryou'll first create it with a list of patterns you'd like to match against and then executematchagainst a node.matcher=MultiPatternMatcher([("rule1",'(Literal #value="re.*ex")'),("rule2",'(IntLiteral #value="1")')])match=matcher.match(node)the result will be eitherNone, meaning no pattern matched, or a tuple of the name of the pattern that matched and the match dict.The match dict will contain a mapping from capture keys to values. Capture keys are the names you can embed within the pattern to store something within the matched node. E.g.:matcher=MultiPatternMatcher([("rule","(ParentType @child_tuple_field =[(Literal) -> first_child, * -> remaining_children])")])this pattern matches any subclass ofParentTypethat has a child fieldchild_tuple_fieldwhich is a sequence, first child of type Literal and then zero or more children of any type and shape.The capture keyfirst_childwill be mapped to the first child of the node that matched the pattern andremaining_childrenwill be mapped to a tuple of the remaining children.Visitors & Transformerspyoak provides 3 classes for visiting and transforming the tree.ASTVisitorA base class for all visitors, and the one to use if you want to visit the tree without changing it.It provides avisitmethod that will call the appropriatevisit_<node_type>method on the visitor object orgeneric_visitif one is not available.Important notes:The class is generic in visit method return type. It is assumed that the type of allvisit_<node_type>methods will match the return type. Unfirtunately, this cannot be type checked by mypy.It is abstract class. As a minimum you must overridegeneric_visitVisitor doesn't visit children by default. It is up to you to callvisiton the children you want to visit.By default, visitor methods are matched by the type of the node, if not found the next type in the MRO is checked, and so on. Going back to our examples abovevisit_Exprwill be triggered for any subclass of Expr, unless a more specific visitor is defined.If you'd like visitor methods to be matched only by strict type equality, change the class variablestrictto True in your visitor class.You can passvalidate=Truewhen inheriting fromASTVisitorlike soclass MyVisitor(ASTVisitor, validate=True)to validate that all visitor methods match the call signature. I.e. if you have a visitor methodvisit_Expr(self, node: IntExpr)andvalidate=Trueis passed, an exception will be raised.ASTTransformVisitorA base class for all visitors that need to transform the tree. It inherits fromASTVisitor, thus all of the notes above also apply to it.UnlikeASTVisitor:It is not generic. Visitor methods return type is pinned toASTNode | NoneIt provides a default implementation ofgeneric_visitthat will transform all children by default, if it detect any changes it will return a new node with the transformed children. Otherwise, it will return the original node.You can use a helper method_transform_childrento transform children. It will return a dictionary suitable for passing todataclasses.replaceto create a new node with the transformed children.An alias tovisitcalledtransformis provided for convenience.Serialization & DeserializationASTNode is itself based onpyoak.serialize.DataClassSerializeMixin, which itself is an extension ofmashumaro'sDataClassDictMixinwith Json, Yaml & MsgPack (de)serialization support.Most of the mashumaro capabilities are available verbatim, includingomitmetadata. There are a couple of differences though:You should useas_dict/as_objinstead of mashumari'sto_dict/from_dict. The former correctly handles the type tag and all additional features belowas_dict/as_obj,to/from_json,to/from_msgpack, andto_from_yamlmethods provide extra argumentserialization_options. There are a couple of built-in ones to serialize the list of child field names, to sort keys, to omit the type tag, and to serialize origin source values in an optimized way. You can also use this to further customize the serialization/deserialization process.For more information about built-in options, checkpyoak.serializeandpyoak.originmodules docstringsDebugging & Reporting Deserialization Errors for Deeply Nested TreesTrees can get pretty deep (we've seen a real-world usage of this library with millions of nodes, hundreds of levels deep).By default, mashumaro will report an error at the top of the tree with the full source as context. This is not only unhelpful but also impossible to work with when your source is 100s of MBs.For this reason, pyoak provides a helperunwrap_invalid_field_exceptioninpyoak.serialize. It will unwrap the exception and provide a tuple with a path to the invalid field and the inner exception that caused the error.CreditsThe (de)serialization code is based on the awesomemashumaroPattern matching definitions and engine runs onlark-parserLinksGitHub:https://github.com/mishamsk/pyoakPyPI:https://pypi.org/project/pyoak/Free software: MIT