Allow deepcopy of fparser tree

.gitignore

@@ -33,6 +33,7 @@
 # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 *.pyc
 *~
+src/**/*.mod
 *.log
 _build/
 htmlcov/

src/fparser/two/Fortran2003.py

@@ -252,19 +252,23 @@ class Program(BlockBase):  # R201
     use_names = ["Program_Unit"]
 
     @show_result
-    def __new__(cls, string):
+    def __new__(cls, string, _deepcopy=False):
         """Wrapper around base class __new__ to catch an internal NoMatchError
         exception and raise it as an external FortranSyntaxError exception.
 
         :param type cls: the class of object to create
         :param string: (source of) Fortran string to parse
         :type string: :py:class:`FortranReaderBase`
+        :param _deepcopy: Flag to signal whether this class is
+            created by a deep copy
+        :type _deepcopy: bool
+
         :raises FortranSyntaxError: if the code is not valid Fortran
 
         """
         # pylint: disable=unused-argument
         try:
-            return Base.__new__(cls, string)
+            return Base.__new__(cls, string, _deepcopy=_deepcopy)
         except NoMatchError:
             # At the moment there is no useful information provided by
             # NoMatchError so we pass on an empty string.
@@ -277,6 +281,18 @@ def __new__(cls, string):
             # provides line number information).
             raise FortranSyntaxError(string, excinfo)
 
+    def __getnewargs__(self):
+        """Method to dictate the values passed to the __new__() method upon
+        unpickling. The method must return a pair (args, kwargs) where
+        args is a tuple of positional arguments and kwargs a dictionary
+        of named arguments for constructing the object. Those will be
+        passed to the __new__() method upon unpickling.
+
+        :return: set of arguments for __new__
+        :rtype: tuple[str, bool]
+        """
+        return (self.string, True)
+
     @staticmethod
     def match(reader):
         """Implements the matching for a Program. Whilst the rule looks like

src/fparser/two/tests/test_parser.py

@@ -36,7 +36,7 @@
 import pytest
 from fparser.two.parser import ParserFactory
 from fparser.common.readfortran import FortranStringReader
-from fparser.two.utils import FortranSyntaxError
+from fparser.two.utils import FortranSyntaxError, StmtBase
 from fparser.two.symbol_table import SYMBOL_TABLES
 from fparser.two import Fortran2003, Fortran2008
 
@@ -84,3 +84,112 @@ class do not affect current calls.
     with pytest.raises(ValueError) as excinfo:
         parser = ParserFactory().create(std="invalid")
         assert "is an invalid standard" in str(excinfo.value)
+
+
+def _cmp_tree_types_rec(
+    node1: Fortran2003.Program, node2: Fortran2003.Program, depth: int = 0
+):
+    """Helper function to recursively check for deepcopied programs
+
+    :param node1: First AST tree to check
+    :type node1: Fortran2003.Program
+    :param node2: Second AST tree to check
+    :type node2: Fortran2003.Program
+    :param depth: Depth useful later on for debugging reasons,
+        defaults to 0
+    :type depth: int, optional
+    """
+
+    # Make sure that both trees are the same
+    assert type(node1) is type(
+        node2
+    ), f"Nodes have different types: '{type(node1)}' and '{type(node2)}"
+
+    if node1 is None:
+        # Just return for None objects
+        return
+
+    if type(node1) is str:
+        # WARNING: Different string objects with the same can have the same id.
+        # Therefore, we can't compare with 'is' or with 'id(.) == id(.)'.
+        # We can just compare the both strings have the same content.
+        # See https://stackoverflow.com/questions/20753364/why-does-creating-multiple-objects-without-naming-them-result-in-them-having-the
+        assert node1 == node2
+        return
+
+    else:
+        # Make sure that we're working on a copy rather than the same object
+        assert node1 is not node2, "Nodes refer to the same object"
+
+    # Continue recursive traversal of ast
+    for child1, child2 in zip(node1.children, node2.children):
+        _cmp_tree_types_rec(child1, child2, depth + 1)
+
+
+_f90_source_test = """
+module andy
+implicit none
+
+   real :: apple = 1.0
+   real, parameter :: pi = 3.14
+
+contains
+   subroutine sergi()
+      print *, "Pi = ", pi
+      print *,  "apple = ", apple
+   end subroutine
+
+end module andy
+
+
+program awesome
+    use andy
+    implicit none
+
+   real :: x
+   integer :: i
+
+   x = 2.2
+   i = 7
+
+   call sergi()
+
+   print *, "apple pie: ", apple, pi
+   print *, "i: ", i
+
+end program awesome
+
+"""
+
+
+def test_deepcopy():
+    """
+    Test that we can deepcopy a parsed fparser tree.
+    """
+
+    parser = ParserFactory().create(std="f2008")
+    reader = FortranStringReader(_f90_source_test)
+    ast = parser(reader)
+
+    import copy
+
+    new_ast = copy.deepcopy(ast)
+
+    _cmp_tree_types_rec(new_ast, ast)
+
+
+def test_pickle():
+    """
+    Test that we can pickle and unpickle a parsed fparser tree.
+    """
+
+    parser = ParserFactory().create(std="f2008")
+    reader = FortranStringReader(_f90_source_test)
+    ast = parser(reader)
+
+    import pickle
+
+    s = pickle.dumps(ast)
+    new_ast = pickle.loads(s)
+
+    _cmp_tree_types_rec(new_ast, ast)

src/fparser/two/utils.py

@@ -409,7 +409,7 @@ def __init__(self, string, parent_cls=None):
         self.parent = None
 
     @show_result
-    def __new__(cls, string, parent_cls=None):
+    def __new__(cls, string, parent_cls=None, _deepcopy=False):
         if parent_cls is None:
             parent_cls = [cls]
         elif cls not in parent_cls:
@@ -418,6 +418,11 @@ def __new__(cls, string, parent_cls=None):
         # Get the class' match method if it has one
         match = getattr(cls, "match", None)
 
+        if _deepcopy:
+            # If this is part of a deep-copy operation (and string is None), simply call
+            # the super method without string
+            return super().__new__(cls)
+
         if (
             isinstance(string, FortranReaderBase)
             and match
@@ -505,6 +510,18 @@ def __new__(cls, string, parent_cls=None):
             errmsg = f"{cls.__name__}: '{string}'"
         raise NoMatchError(errmsg)
 
+    def __getnewargs__(self):
+        """Method to dictate the values passed to the __new__() method upon
+        unpickling. The method must return a pair (args, kwargs) where
+        args is a tuple of positional arguments and kwargs a dictionary
+        of named arguments for constructing the object. Those will be
+        passed to the __new__() method upon unpickling.
+
+        :return: set of arguments for __new__
+        :rtype: tuple[str, NoneType, bool]
+        """
+        return (self.string, None, True)
+
     def get_root(self):
         """
         Gets the node at the root of the parse tree to which this node belongs.