Implement a fridge-copy-dir recipe

scone/default/recipes/fridge.py

@@ -20,20 +20,27 @@ import logging
 import os
 from asyncio import Future
 from pathlib import Path
-from typing import Dict, cast
+from typing import Dict, List, Set, Tuple, cast
 from urllib.parse import urlparse
 
 import requests
 
 from scone.common.misc import sha256_file
-from scone.common.modeutils import DEFAULT_MODE_FILE, parse_mode
+from scone.common.modeutils import DEFAULT_MODE_DIR, DEFAULT_MODE_FILE, parse_mode
 from scone.default.steps import fridge_steps
 from scone.default.steps.fridge_steps import (
     SUPERMARKET_RELATIVE,
     FridgeMetadata,
     load_and_transform,
 )
-from scone.default.utensils.basic_utensils import Chmod, Chown, HashFile, WriteFile
+from scone.default.utensils.basic_utensils import (
+    Chmod,
+    Chown,
+    HashFile,
+    MakeDirectory,
+    Stat,
+    WriteFile,
+)
 from scone.head.head import Head
 from scone.head.kitchen import Kitchen, Preparation
 from scone.head.recipe import Recipe, RecipeContext
@@ -103,6 +110,117 @@ class FridgeCopy(Recipe):
         k.get_dependency_tracker().register_fridge_file(self._desugared_src)
 
 
+class FridgeCopyDir(Recipe):
+    """
+    Declares that a directory(!) should be copied from the head to the sous,
+    and optionally, remote files deleted.
+    """
+
+    _NAME = "fridge-copy-dir"
+
+    def __init__(self, recipe_context: RecipeContext, args: dict, head: Head):
+        super().__init__(recipe_context, args, head)
+
+        files = fridge_steps.search_children_in_fridge(head, args["src"])
+        if not files:
+            raise ValueError(
+                f"Cannot find children of directory {args['src']}"
+                f" in the fridge (empty directories not allowed)."
+            )
+
+        self.files: List[Tuple[str, str, Path]] = files
+
+        dest = check_type(args["dest"], str)
+
+        self.dest_dir = Path(dest)
+
+        self.destinations: List[Path] = []
+
+        self.mkdirs: Set[str] = set()
+
+        for relative, relative_unprefixed, full_path in self.files:
+            unextended_path_str, _ = fridge_steps.decode_fridge_extension(
+                relative_unprefixed
+            )
+            self.destinations.append(Path(args["dest"], unextended_path_str))
+            pieces = relative_unprefixed.split("/")
+            for end_index in range(0, len(pieces)):
+                self.mkdirs.add("/".join(pieces[0:end_index]))
+
+        mode = args.get("mode", DEFAULT_MODE_FILE)
+        dir_mode = args.get("mode_dir", args.get("mode", DEFAULT_MODE_DIR))
+        assert isinstance(mode, str) or isinstance(mode, int)
+        assert isinstance(dir_mode, str) or isinstance(dir_mode, int)
+
+        self.file_mode = parse_mode(mode, directory=False)
+        self.dir_mode = parse_mode(dir_mode, directory=True)
+
+    def prepare(self, preparation: Preparation, head: Head) -> None:
+        super().prepare(preparation, head)
+
+        preparation.needs("directory", str(self.dest_dir.parent))
+
+        for mkdir in self.mkdirs:
+            preparation.provides("directory", str(Path(self.dest_dir, mkdir)))
+
+        for (_relative, relative_unprefixed, full_path), destination in zip(
+            self.files, self.destinations
+        ):
+            unextended_path_str, _ = fridge_steps.decode_fridge_extension(
+                relative_unprefixed
+            )
+            preparation.provides("file", str(destination))
+
+    async def cook(self, k: Kitchen) -> None:
+        # create all needed directories
+        for mkdir in self.mkdirs:
+            directory = str(Path(self.dest_dir, mkdir))
+            # print("mkdir ", directory)
+
+            stat = await k.ut1a(Stat(directory), Stat.Result)
+            if stat is None:
+                # doesn't exist, make it
+                await k.ut0(MakeDirectory(directory, self.dir_mode))
+
+            stat = await k.ut1a(Stat(directory), Stat.Result)
+            if stat is None:
+                raise RuntimeError("Directory vanished after creation!")
+
+            if stat.dir:
+                # if (stat.user, stat.group) != (self.targ_user, self.targ_group):
+                #     # need to chown
+                #     await k.ut0(Chown(directory, self.targ_user, self.targ_group))
+
+                if stat.mode != self.dir_mode:
+                    await k.ut0(Chmod(directory, self.dir_mode))
+            else:
+                raise RuntimeError("Already exists but not a dir: " + directory)
+
+        # copy all files from the fridge
+        for (relative, relative_unprefixed, full_local_path), destination in zip(
+            self.files, self.destinations
+        ):
+            unextended_path_str, meta = fridge_steps.decode_fridge_extension(
+                relative_unprefixed
+            )
+            full_remote_path = str(Path(self.dest_dir, unextended_path_str))
+            # print("fcp ", relative, " → ", full_remote_path)
+
+            data = await load_and_transform(
+                k, meta, full_local_path, self.recipe_context.sous
+            )
+            dest_str = str(full_remote_path)
+            chan = await k.start(WriteFile(dest_str, self.file_mode))
+            await chan.send(data)
+            await chan.send(None)
+            if await chan.recv() != "OK":
+                raise RuntimeError(
+                    f"WriteFail failed on fridge-copy to {full_remote_path}"
+                )
+
+            k.get_dependency_tracker().register_fridge_file(relative)
+
+
 class Supermarket(Recipe):
     """
     Downloads an asset (cached if necessary) and copies to sous.

scone/default/steps/fridge_steps.py

@@ -14,7 +14,7 @@
 #
 #  You should have received a copy of the GNU General Public License
 #  along with Scone.  If not, see <https://www.gnu.org/licenses/>.
-
+import os
 from enum import Enum
 from pathlib import Path, PurePath
 from typing import List, Optional, Tuple, Union
@@ -114,3 +114,57 @@ async def load_and_transform(
         #     template.environment.handle_exception()
 
     return data
+
+
+def _find_files_in_dir(relative: str, dir: Path) -> List[Tuple[str, str, Path]]:
+    """
+    :param relative:
+    :param dir:
+    :return: Tuple of (
+        relative path with prefix included,
+        relative path with prefix not included,
+        path to local file
+    )
+    """
+    result = []
+    num_prefix_parts = len(dir.parts)
+    for root, dirs, files in os.walk(dir):
+        for file in files:
+            full_path = Path(root, file)
+            parts = full_path.parts
+            if parts[0:num_prefix_parts] != dir.parts:
+                raise RuntimeError(f"{parts[0:num_prefix_parts]!r} != {dir.parts!r}")
+            dir_relative_path = "/".join(parts[num_prefix_parts:])
+            result.append(
+                (relative + "/" + dir_relative_path, dir_relative_path, full_path)
+            )
+    return result
+
+
+def search_children_in_fridge(
+    head: Head, relative: Union[str, PurePath]
+) -> Optional[List[Tuple[str, str, Path]]]:
+    """
+    Similar to `search_in_fridge` but finds (recursively) ALL children of a named
+    directory. This 'directory' can be split across multiple fridge search paths.
+    """
+    fridge_dirs = get_fridge_dirs(head)
+
+    results = []
+    # only the first file found for a path counts — this allows overrides
+    found_filenames = set()
+
+    for directory in fridge_dirs:
+        potential_path = directory.joinpath(relative)
+        if potential_path.exists():
+            # find children
+            for rel, rel_unprefixed, file in _find_files_in_dir(
+                str(relative), potential_path
+            ):
+                unextended_name, _transformer = decode_fridge_extension(rel)
+                if unextended_name in found_filenames:
+                    continue
+                results.append((rel, rel_unprefixed, file))
+                found_filenames.add(unextended_name)
+
+    return results