Add localization support for directory targets.

law/target/file.py

@@ -289,6 +289,10 @@ def move_to_local(self, *args, **kwargs):
     def move_from_local(self, *args, **kwargs):
         return
 
+    @abstractmethod
+    @contextmanager
+    def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
+        return
 
 class FileSystemFileTarget(FileSystemTarget):
 
@@ -337,11 +341,6 @@ def move_from(self, src, perm=None, dir_perm=None, **kwargs):
         # TODO: complain when src not local? forward to copy_to request depending on protocol?
         return self.fs.move(get_path(src), self.path, perm=perm, dir_perm=dir_perm, **kwargs)
 
-    @abstractmethod
-    @contextmanager
-    def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
-        return
-
 
 class FileSystemDirectoryTarget(FileSystemTarget):
 
@@ -410,9 +409,13 @@ def copy_to(self, dst, perm=None, dir_perm=None, **kwargs):
                     t = self.child(basename, type=type_flag)
                     t.copy_to(os.path.join(_dst, basename), perm=perm, dir_perm=dir_perm, **kwargs)
 
+        return _dst
+
     def copy_from(self, src, perm=None, dir_perm=None, **kwargs):
+        # when src is a directory target itself, forward to its copy_to implementation as it might
+        # be more performant to use its own directory walking
         if isinstance(src, FileSystemDirectoryTarget):
-            return src.copy_to(self.path, perm=perm, dir_perm=dir_perm, **kwargs)
+            return src.copy_to(self, perm=perm, dir_perm=dir_perm, **kwargs)
 
         # create the target dir
         self.touch(perm=dir_perm, **kwargs)
@@ -428,6 +431,8 @@ def copy_from(self, src, perm=None, dir_perm=None, **kwargs):
                     t = self.child(basename, type=type_flag)
                     t.copy_from(os.path.join(_src, basename), perm=perm, dir_perm=dir_perm, **kwargs)
 
+        return self.abspath
+
     def move_to(self, dst, perm=None, dir_perm=None, **kwargs):
         # create the target dir
         _dst = get_path(dst)
@@ -448,9 +453,13 @@ def move_to(self, dst, perm=None, dir_perm=None, **kwargs):
         # finally remove
         self.remove()
 
+        return _dst
+
     def move_from(self, src, perm=None, dir_perm=None, **kwargs):
+        # when src is a directory target itself, forward to its move_to implementation as it might
+        # be more performant to use its own directory walking
         if isinstance(src, FileSystemDirectoryTarget):
-            return src.move_to(self.path, perm=perm, dir_perm=dir_perm, **kwargs)
+            return src.move_to(self, perm=perm, dir_perm=dir_perm, **kwargs)
 
         # create the target dir
         self.touch(perm=dir_perm, **kwargs)
@@ -469,6 +478,8 @@ def move_from(self, src, perm=None, dir_perm=None, **kwargs):
         # finally remove
         self.fs.remove(_src)
 
+        return self.abspath
+
 
 FileSystemTarget.file_class = FileSystemFileTarget
 FileSystemTarget.directory_class = FileSystemDirectoryTarget

law/target/local.py

@@ -425,7 +425,7 @@ def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
         if mode not in ["r", "w", "a"]:
             raise Exception("unknown mode '{}', use 'r', 'w' or 'a'".format(mode))
 
-        logger.debug("localizing file target {!r} with mode '{}'".format(self, mode))
+        logger.debug("localizing {!r} with mode '{}'".format(self, mode))
 
         # get additional arguments
         is_tmp = kwargs.pop("is_tmp", mode in ("w", "a"))
@@ -462,10 +462,10 @@ def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
 
                     # move back again
                     if tmp.exists():
-                        tmp.move_to_local(self, perm=perm, dir_perm=dir_perm)
+                        tmp.copy_to_local(self, perm=perm, dir_perm=dir_perm)
                     else:
-                        logger.warning("cannot move non-existing, temporary target to localized "
-                            "file target {!r}".format(self))
+                        logger.warning("cannot move non-existing localized target to actual "
+                            "representation {!r}".format(self))
                 finally:
                     tmp.remove()
             else:
@@ -485,6 +485,66 @@ def _child_args(self, path):
         kwargs["fs"] = self.fs
         return args, kwargs
 
+    @contextmanager
+    def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
+        if mode not in ["r", "w", "a"]:
+            raise Exception("unknown mode '{}', use 'r', 'w' or 'a'".format(mode))
+
+        logger.debug("localizing {!r} with mode '{}'".format(self, mode))
+
+        # get additional arguments
+        is_tmp = kwargs.pop("is_tmp", mode in ("w", "a"))
+
+        if mode == "r":
+            if is_tmp:
+                # create a temporary target
+                tmp = self.__class__(is_tmp=True, tmp_dir=tmp_dir)
+
+                # copy contents
+                self.copy_to_local(tmp)
+
+                # yield the copy
+                try:
+                    yield tmp
+                finally:
+                    tmp.remove(silent=True)
+            else:
+                # simply yield
+                yield self
+
+        else:  # mode "w" or "a"
+            if is_tmp:
+                # create a temporary target
+                tmp = self.__class__(is_tmp=True, tmp_dir=tmp_dir)
+
+                # copy in append mode, otherwise ensure that it exists
+                if mode == "a" and self.exists():
+                    self.copy_to_local(tmp)
+                else:
+                    tmp.touch()
+
+                # yield the copy
+                try:
+                    yield tmp
+
+                    # move back again, first removing current content
+                    # TODO: keep track of changed contents in "a" mode and copy only those?
+                    if tmp.exists():
+                        self.remove()
+                        tmp.copy_to_local(self, perm=perm, dir_perm=dir_perm)
+                    else:
+                        logger.warning("cannot move non-existing localized target to actual "
+                            "representation {!r}, leaving original contents unchanged".format(self))
+                finally:
+                    tmp.remove()
+            else:
+                # create the parent dir and the directory itself
+                self.parent.touch(perm=dir_perm)
+                self.touch(perm=perm)
+
+                # simply yield, do not differentiate "w" and "a" modes
+                yield self
+
 
 LocalTarget.file_class = LocalFileTarget
 LocalTarget.directory_class = LocalDirectoryTarget

law/target/remote/base.py

@@ -19,7 +19,7 @@
     FileSystem, FileSystemTarget, FileSystemFileTarget, FileSystemDirectoryTarget, get_path,
     get_scheme, add_scheme, remove_scheme,
 )
-from law.target.local import LocalFileSystem, LocalFileTarget
+from law.target.local import LocalFileSystem, LocalFileTarget, LocalDirectoryTarget
 from law.target.remote.cache import RemoteCache
 from law.target.formatter import find_formatter
 from law.util import make_list, merge_dicts
@@ -649,7 +649,7 @@ def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
         if mode not in ["r", "w", "a"]:
             raise Exception("unknown mode '{}', use 'r', 'w' or 'a'".format(mode))
 
-        logger.debug("localizing file target {!r} with mode '{}'".format(self, mode))
+        logger.debug("localizing {!r} with mode '{}'".format(self, mode))
 
         if mode == "r":
             with self.fs.open(self.path, "r", _yield_path=True, perm=perm, **kwargs) as lpath:
@@ -660,16 +660,16 @@ def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
 
             # copy to local in append mode
             if mode == "a" and self.exists():
-                self.copy_to_local(tmp)
+                self.copy_to_local(tmp, **kwargs)
 
             try:
                 yield tmp
 
                 if tmp.exists():
                     self.copy_from_local(tmp, perm=perm, dir_perm=dir_perm, **kwargs)
                 else:
-                    logger.warning("cannot move non-existing localized file target {!r}".format(
-                        self))
+                    logger.warning("cannot move non-existing localized target to actual "
+                        "representation {!r}".format(self))
             finally:
                 tmp.remove()
 
@@ -681,6 +681,51 @@ def _child_args(self, path):
         args += (self.fs,)
         return args, kwargs
 
+    @contextmanager
+    def localize(self, mode="r", perm=None, dir_perm=None, tmp_dir=None, **kwargs):
+        if mode not in ["r", "w", "a"]:
+            raise Exception("unknown mode '{}', use 'r', 'w' or 'a'".format(mode))
+
+        logger.debug("localizing {!r} with mode '{}'".format(self, mode))
+
+        if mode == "r":
+            # create a temporary directory
+            tmp = LocalDirectoryTarget(is_tmp=True, tmp_dir=tmp_dir)
+
+            # copy contents
+            self.copy_to_local(tmp, **kwargs)
+
+            # yield the copy
+            try:
+                yield tmp
+            finally:
+                tmp.remove(silent=True)
+
+        else:  # mode "w" or "a"
+            # create a temporary directory
+            tmp = LocalDirectoryTarget(is_tmp=True, tmp_dir=tmp_dir)
+
+            # copy in append mode, otherwise ensure that it exists
+            if mode == "a" and self.exists():
+                self.copy_to_local(tmp)
+            else:
+                tmp.touch()
+
+            # yield the copy
+            try:
+                yield tmp
+
+                # move back again, first removing current content
+                # TODO: keep track of changed contents in "a" mode and copy only those?
+                if tmp.exists():
+                    self.remove(**kwargs)
+                    self.copy_from_local(tmp, perm=perm, dir_perm=dir_perm, **kwargs)
+                else:
+                    logger.warning("cannot move non-existing localized target to actual "
+                        "representation {!r}, leaving original contents unchanged".format(self))
+            finally:
+                tmp.remove()
+
 
 RemoteTarget.file_class = RemoteFileTarget
 RemoteTarget.directory_class = RemoteDirectoryTarget