Merge pull request #127 from stfc/add_get_child

Add get_child() utility

CHANGELOG.md

@@ -11,6 +11,8 @@ Modifications by (in alphabetical order):
 * A. R. Porter, Science & Technology Facilities Council, UK
 * P. Vitt, University of Siegen, Germany
 
+21/11/2018 PR #127 for #126. Adds get_child function to help AST traversal.
+
 19/11/2018 PR #124 for #112. Bug fix - spaces within names are now
            rejected by fparser2.
 

doc/fparser2.rst

@@ -130,7 +130,7 @@ executing the AST. For example:
 
 Note that the two readers will ignore (and dispose of) comments by
 default. If you wish comments to be retained then you must set
-`ignore_comments=True` when creating the reader. The AST created by
+`ignore_comments=False` when creating the reader. The AST created by
 fparser2 will then have `Comment` nodes representing any comments
 found in the code. Nodes representing in-line comments will be added
 immediately following the node representing the code in which they
@@ -183,3 +183,15 @@ tree representing the parsed code are instances of either `BlockBase`
 or `SequenceBase`. Child nodes are then stored in the `.content`
 attribute of `BlockBase` objects or the `.items` attribute of
 `SequenceBase` objects. Both of these attributes are Tuple instances.
+
+
+Walking the AST
+---------------
+
+fparser2 provides two functions to support the traversal of the
+AST that it constructs:
+
+.. automethod:: fparser.two.utils.walk_ast
+
+.. automethod:: fparser.two.utils.get_child
+

src/fparser/two/tests/test_utils.py

@@ -110,3 +110,29 @@ def test_blockbase_match_name_classes(f2003_create):
         ast = If_Construct(reader)
     assert ("at line 2\n>>>endif label\nName 'label' has no corresponding "
             "starting name") in str(excinfo.value)
+
+
+def test_get_child(f2003_create):
+    ''' Test the get_child() utility. '''
+    from fparser.two import Fortran2003
+    from fparser.two.utils import get_child, walk_ast
+    reader = get_reader("program hello\n"
+                        "write(*,*) 'hello'\n"
+                        "write(*,*) 'goodbye'\n"
+                        "end program hello\n")
+    main = Fortran2003.Program(reader)
+    prog = get_child(main, Fortran2003.Main_Program)
+    exe = get_child(prog, Fortran2003.Execution_Part)
+    assert isinstance(exe, Fortran2003.Execution_Part)
+    write_stmt = get_child(exe, Fortran2003.Write_Stmt)
+    # Check that we got the first write and not the second
+    assert "goodbye" not in str(write_stmt)
+    # The top level has no Io_Control_Spec children
+    assert not get_child(main, Fortran2003.Io_Control_Spec)
+    # Check functionality when node has children in `items` and
+    # not in `content`
+    io_nodes = walk_ast(main.content, my_types=[Fortran2003.Io_Control_Spec])
+    assert not hasattr(io_nodes[0], "content")
+    io_unit = get_child(io_nodes[0], Fortran2003.Io_Unit)
+    assert isinstance(io_unit, Fortran2003.Io_Unit)
+

src/fparser/two/utils.py

@@ -198,9 +198,9 @@ class Base(ComparableMixin):
     ''' Base class for Fortran 2003 syntax rules.
 
     All Base classes have the following attributes:
-      \.string - original argument to construct a class instance, its type \
-                is either str or FortranReaderBase.
-      \.item   - Line instance (holds label) or None.
+      self.string - original argument to construct a class instance, its type \
+                    is either str or FortranReaderBase.
+      self.item   - Line instance (holds label) or None.
 
     '''
     # This dict of subclasses is populated dynamically by code at the end
@@ -1269,9 +1269,22 @@ def tostr(self):
 
 
 def walk_ast(children, my_types=None, indent=0, debug=False):
-    '''' Walk down the tree produced by fparser2 where children
+    '''
+    Walk down the tree produced by fparser2 where children
     are listed under 'content'.  Returns a list of all nodes with the
-    specified type(s). '''
+    specified type(s).
+
+    :param children: list of child nodes from which to walk.
+    :type children: list of :py:class:fparser.two.utils.Base.
+    :param my_types: list of types of Node to return. (Default is to \
+                     return all nodes.)
+    :type my_types: list of type
+    :param int indent: extent to which to indent debug output.
+    :param bool debug: whether or not to write textual representation of AST \
+                       to stdout.
+    :returns: a list of nodes
+    :rtype: `list` of :py:class:`fparser.two.utils.Base`
+    '''
     local_list = []
     for child in children:
         if debug:
@@ -1292,3 +1305,27 @@ def walk_ast(children, my_types=None, indent=0, debug=False):
             local_list += walk_ast(child.items, my_types, indent+1, debug)
 
     return local_list
+
+
+def get_child(root_node, node_type):
+    '''
+    Searches for the first immediate child of root_node that is of the
+    specified type.
+
+    :param root_node: the parent of the child nodes we will search through.
+    :type root_node: :py:class:`fparser.two.utils.Base`
+    :param type node_type: the class of child node to search for.
+
+    :returns: the first child node of type node_type that is encountered or \
+              None.
+    :rtype: :py:class:`fparser.two.utils.Base`
+    '''
+    children = []
+    if hasattr(root_node, "content"):
+        children = root_node.content
+    elif hasattr(root_node, "items"):
+        children = root_node.items
+    for node in children:
+        if isinstance(node, node_type):
+            return node
+    return None