Add base class for directory visitor (#32008)

3 files changed, 103 insertions, 31 deletions

diff --git a/lib/spack/llnl/util/filesystem.py b/lib/spack/llnl/util/filesystem.py
index 1740fb71c0..a690ef3c37 100644
--- a/lib/spack/llnl/util/filesystem.py
+++ b/lib/spack/llnl/util/filesystem.py
@@ -81,6 +81,8 @@ __all__ = [
     "unset_executable_mode",
     "working_dir",
     "keep_modification_time",
+    "BaseDirectoryVisitor",
+    "visit_directory_tree",
 ]
 
 
@@ -1133,20 +1135,89 @@ def lexists_islink_isdir(path):
     return True, is_link, is_dir
 
 
+class BaseDirectoryVisitor(object):
+    """Base class and interface for :py:func:`visit_directory_tree`."""
+
+    def visit_file(self, root, rel_path, depth):
+        """Handle the non-symlink file at ``os.path.join(root, rel_path)``
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current file from ``root``
+            depth (int): depth of current file from the ``root`` directory"""
+        pass
+
+    def visit_symlinked_file(self, root, rel_path, depth):
+        """Handle the symlink to a file at ``os.path.join(root, rel_path)``.
+        Note: ``rel_path`` is the location of the symlink, not to what it is
+        pointing to. The symlink may be dangling.
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current symlink from ``root``
+            depth (int): depth of current symlink from the ``root`` directory"""
+        pass
+
+    def before_visit_dir(self, root, rel_path, depth):
+        """Return True from this function to recurse into the directory at
+        os.path.join(root, rel_path). Return False in order not to recurse further.
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current directory from ``root``
+            depth (int): depth of current directory from the ``root`` directory
+
+        Returns:
+            bool: ``True`` when the directory should be recursed into. ``False`` when
+            not"""
+        return False
+
+    def before_visit_symlinked_dir(self, root, rel_path, depth):
+        """Return ``True`` to recurse into the symlinked directory and ``False`` in
+        order not to. Note: ``rel_path`` is the path to the symlink itself.
+        Following symlinked directories blindly can cause infinite recursion due to
+        cycles.
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current symlink from ``root``
+            depth (int): depth of current symlink from the ``root`` directory
+
+        Returns:
+            bool: ``True`` when the directory should be recursed into. ``False`` when
+            not"""
+        return False
+
+    def after_visit_dir(self, root, rel_path, depth):
+        """Called after recursion into ``rel_path`` finished. This function is not
+        called when ``rel_path`` was not recursed into.
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current directory from ``root``
+            depth (int): depth of current directory from the ``root`` directory"""
+        pass
+
+    def after_visit_symlinked_dir(self, root, rel_path, depth):
+        """Called after recursion into ``rel_path`` finished. This function is not
+        called when ``rel_path`` was not recursed into.
+
+        Parameters:
+            root (str): root directory
+            rel_path (str): relative path to current symlink from ``root``
+            depth (int): depth of current symlink from the ``root`` directory"""
+        pass
+
+
 def visit_directory_tree(root, visitor, rel_path="", depth=0):
-    """
-    Recurses the directory root depth-first through a visitor pattern
-
-    The visitor interface is as follows:
-    - visit_file(root, rel_path, depth)
-    - before_visit_dir(root, rel_path, depth) -> bool
-        if True, descends into this directory
-    - before_visit_symlinked_dir(root, rel_path, depth) -> bool
-        if True, descends into this directory
-    - after_visit_dir(root, rel_path, depth) -> void
-        only called when before_visit_dir returns True
-    - after_visit_symlinked_dir(root, rel_path, depth) -> void
-        only called when before_visit_symlinked_dir returns True
+    """Recurses the directory root depth-first through a visitor pattern using the
+    interface from :py:class:`BaseDirectoryVisitor`
+
+    Parameters:
+        root (str): path of directory to recurse into
+        visitor (BaseDirectoryVisitor): what visitor to use
+        rel_path (str): current relative path from the root
+        depth (str): current depth from the root
     """
     dir = os.path.join(root, rel_path)
 
@@ -1190,9 +1261,11 @@ def visit_directory_tree(root, visitor, rel_path="", depth=0):
             if not lexists:
                 continue
 
-        if not isdir:
-            # handle files
+        if not isdir and not islink:
+            # handle non-symlink files
             visitor.visit_file(root, rel_child, depth)
+        elif not isdir:
+            visitor.visit_symlinked_file(root, rel_child, depth)
         elif not islink and visitor.before_visit_dir(root, rel_child, depth):
             # Handle ordinary directories
             visit_directory_tree(root, visitor, rel_child, depth + 1)
diff --git a/lib/spack/llnl/util/link_tree.py b/lib/spack/llnl/util/link_tree.py
index 947ca9c541..b29a18c979 100644
--- a/lib/spack/llnl/util/link_tree.py
+++ b/lib/spack/llnl/util/link_tree.py
@@ -13,7 +13,7 @@ import shutil
 from collections import OrderedDict
 
 import llnl.util.tty as tty
-from llnl.util.filesystem import mkdirp, touch, traverse_tree
+from llnl.util.filesystem import BaseDirectoryVisitor, mkdirp, touch, traverse_tree
 from llnl.util.symlink import islink, symlink
 
 __all__ = ["LinkTree"]
@@ -45,7 +45,7 @@ class MergeConflict:
         self.src_b = src_b
 
 
-class SourceMergeVisitor(object):
+class SourceMergeVisitor(BaseDirectoryVisitor):
     """
     Visitor that produces actions:
     - An ordered list of directories to create in dst
@@ -106,9 +106,6 @@ class SourceMergeVisitor(object):
             self.directories[proj_rel_path] = (root, rel_path)
             return True
 
-    def after_visit_dir(self, root, rel_path, depth):
-        pass
-
     def before_visit_symlinked_dir(self, root, rel_path, depth):
         """
         Replace symlinked dirs with actual directories when possible in low depths,
@@ -141,9 +138,6 @@ class SourceMergeVisitor(object):
         self.visit_file(root, rel_path, depth)
         return False
 
-    def after_visit_symlinked_dir(self, root, rel_path, depth):
-        pass
-
     def visit_file(self, root, rel_path, depth):
         proj_rel_path = os.path.join(self.projection, rel_path)
 
@@ -173,6 +167,10 @@ class SourceMergeVisitor(object):
             # Otherwise register this file to be linked.
             self.files[proj_rel_path] = (root, rel_path)
 
+    def visit_symlinked_file(self, root, rel_path, depth):
+        # Treat symlinked files as ordinary files (without "dereferencing")
+        self.visit_file(root, rel_path, depth)
+
     def set_projection(self, projection):
         self.projection = os.path.normpath(projection)
 
@@ -200,7 +198,7 @@ class SourceMergeVisitor(object):
                 )
 
 
-class DestinationMergeVisitor(object):
+class DestinationMergeVisitor(BaseDirectoryVisitor):
     """DestinatinoMergeVisitor takes a SourceMergeVisitor
     and:
 
@@ -240,9 +238,6 @@ class DestinationMergeVisitor(object):
         # don't descend into it.
         return False
 
-    def after_visit_dir(self, root, rel_path, depth):
-        pass
-
     def before_visit_symlinked_dir(self, root, rel_path, depth):
         """
         Symlinked directories in the destination prefix should
@@ -269,9 +264,6 @@ class DestinationMergeVisitor(object):
         # Never descend into symlinked target dirs.
         return False
 
-    def after_visit_symlinked_dir(self, root, rel_path, depth):
-        pass
-
     def visit_file(self, root, rel_path, depth):
         # Can't merge a file if target already exists
         if rel_path in self.src.directories:
@@ -290,6 +282,10 @@ class DestinationMergeVisitor(object):
                 )
             )
 
+    def visit_symlinked_file(self, root, rel_path, depth):
+        # Treat symlinked files as ordinary files (without "dereferencing")
+        self.visit_file(root, rel_path, depth)
+
 
 class LinkTree(object):
     """Class to create trees of symbolic links from a source directory.
diff --git a/lib/spack/spack/test/llnl/util/filesystem.py b/lib/spack/spack/test/llnl/util/filesystem.py
index 559ff5bace..595142f878 100644
--- a/lib/spack/spack/test/llnl/util/filesystem.py
+++ b/lib/spack/spack/test/llnl/util/filesystem.py
@@ -729,7 +729,7 @@ def test_lexists_islink_isdir(tmpdir):
     assert fs.lexists_islink_isdir(symlink_to_symlink_to_file) == (True, True, False)
 
 
-class RegisterVisitor(object):
+class RegisterVisitor(fs.BaseDirectoryVisitor):
     """A directory visitor that keeps track of all visited paths"""
 
     def __init__(self, root, follow_dirs=True, follow_symlink_dirs=True):
@@ -751,6 +751,9 @@ class RegisterVisitor(object):
         self.check(root, rel_path, depth)
         self.files.append(rel_path)
 
+    def visit_symlinked_file(self, root, rel_path, depth):
+        self.visit_file(root, rel_path, depth)
+
     def before_visit_dir(self, root, rel_path, depth):
         self.check(root, rel_path, depth)
         self.dirs_before.append(rel_path)