Use Path objects in AssetManager

grizzly/common/storage.py

@@ -175,7 +175,7 @@ def clone(self):
         result.assets = dict(self.assets)
         if result.assets:
             assert self.assets_path
-            org_path = Path(self.assets_path)
+            org_path = self.assets_path
             try:
                 # copy asset data from test case
                 result.assets_path = result.root / org_path.relative_to(self.root)
@@ -273,7 +273,7 @@ def dump(self, dst_path, include_details=False):
             # save target assets and update meta data
             if self.assets:
                 assert isinstance(self.assets, dict)
-                assert isinstance(self.assets_path, (Path, str))
+                assert isinstance(self.assets_path, Path)
                 info["assets"] = self.assets
                 info["assets_path"] = "_assets_"
                 copytree(self.assets_path, dst_path / info["assets_path"])
@@ -439,7 +439,7 @@ def load_single(cls, path, adjacent=False, copy=True):
         )
         # load all adjacent data from directory
         if adjacent:
-            for entry in Path(entry_point.parent).rglob("*"):
+            for entry in entry_point.parent.rglob("*"):
                 if not entry.is_file():
                     continue
                 # ignore asset path

grizzly/common/test_storage.py

@@ -5,7 +5,6 @@
 
 from itertools import chain
 from json import dumps, loads
-from pathlib import Path
 from zipfile import ZIP_DEFLATED, ZipFile
 
 from pytest import mark, raises
@@ -219,8 +218,8 @@ def test_testcase_08(tmp_path):
     (tmp_path / "src" / "nested" / "empty").mkdir()
     dst_dir = tmp_path / "dst"
     # build test case
-    with AssetManager(base_path=str(tmp_path)) as assets:
-        assets.add("example", str(asset_file))
+    with AssetManager(base_path=tmp_path) as assets:
+        assets.add("example", asset_file)
         with TestCase("target.bin", "test-adapter") as src:
             src.env_vars["TEST_ENV_VAR"] = "100"
             src.add_from_file(entry_point)
@@ -238,7 +237,7 @@ def test_testcase_08(tmp_path):
     with TestCase.load_single(dst_dir) as dst:
         assert "example" in dst.assets
         assert dst.assets_path is not None
-        assert (Path(dst.assets_path) / "asset.bin").is_file()
+        assert (dst.assets_path / "asset.bin").is_file()
         assert "_assets_/asset.bin" not in (x.file_name for x in dst._files.optional)
         assert dst.entry_point == "target.bin"
         assert "target.bin" in (x.file_name for x in dst._files.required)
@@ -293,7 +292,7 @@ def test_testcase_10(tmp_path):
         org.time_limit = 10
         org.add_from_bytes(b"a", "a.html")
         org.assets = {"sample": asset.name}
-        org.assets_path = str(asset_path)
+        org.assets_path = asset_path
         org.dump(working, include_details=True)
         assert (working / "_assets_" / asset.name).is_file()
         with TestCase.load_single(working, adjacent=False) as loaded:
@@ -335,8 +334,8 @@ def test_testcase_14(tmp_path):
     nested.mkdir()
     asset_file = tmp_path / "example_asset"
     asset_file.touch()
-    with AssetManager(base_path=str(tmp_path)) as assets:
-        assets.add("example", str(asset_file))
+    with AssetManager(base_path=tmp_path) as assets:
+        assets.add("example", asset_file)
         with TestCase("target.bin", "test-adapter") as src:
             src.assets = assets.assets
             src.assets_path = assets.path

grizzly/reduce/test_reduce.py

@@ -857,7 +857,7 @@ def submit(test_cases, report, force=False):
     target.filtered_environ.return_value = {"test": "abc"}
     with AssetManager(base_path=str(tmp_path)) as assets:
         (tmp_path / "example_asset").touch()
-        assets.add("example", str(tmp_path / "example_asset"), copy=False)
+        assets.add("example", tmp_path / "example_asset", copy=False)
         target.assets = assets
         try:
             mgr = ReduceManager(

grizzly/replay/replay.py

@@ -179,9 +179,7 @@ def load_testcases(cls, path, subset=None):
         env_vars = None
         for test in testcases:
             if assets is None and test.assets and test.assets_path:
-                assets = AssetManager.load(
-                    test.assets, str(test.root / test.assets_path)
-                )
+                assets = AssetManager.load(test.assets, test.root / test.assets_path)
             if not env_vars and test.env_vars:
                 env_vars = dict(test.env_vars)
             test.env_vars.clear()
@@ -658,8 +656,7 @@ def main(cls, args):
                 LOG.debug("adding environment loaded from test case")
                 target.merge_environment(env_vars)
 
-            # TODO: support overriding existing assets
-            # prioritize specified assets over included
+            # TODO: prioritize specified assets over included
             target.assets.add_batch(args.asset)
             target.process_assets()
 

grizzly/replay/test_replay.py

@@ -637,7 +637,7 @@ def test_replay_21(tmp_path):
     # build test case
     with AssetManager() as assets:
         (tmp_path / "prefs.js").touch()
-        assets.add("prefs", str(tmp_path / "prefs.js"), copy=False)
+        assets.add("prefs", tmp_path / "prefs.js", copy=False)
         with TestCase(data.name, "foo") as src:
             src.env_vars = {"foo": "bar"}
             src.assets = dict(assets.assets)

grizzly/target/assets.py

@@ -2,9 +2,7 @@
 # License, v. 2.0. If a copy of the MPL was not distributed with this
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 from logging import getLogger
-from os import makedirs, unlink
-from os.path import abspath, basename, exists, isdir, isfile
-from os.path import join as pathjoin
+from pathlib import Path
 from shutil import copyfile, copytree, move, rmtree
 from tempfile import mkdtemp
 
@@ -19,13 +17,12 @@ class AssetError(Exception):
     """Raised by AssetManager"""
 
 
-# TODO: use pathlib
 class AssetManager:
     __slots__ = ("assets", "path")
 
     def __init__(self, base_path=None):
         self.assets = {}
-        self.path = mkdtemp(prefix="assets_", dir=base_path)
+        self.path = Path(mkdtemp(prefix="assets_", dir=base_path))
 
     def __enter__(self):
         return self
@@ -38,41 +35,36 @@ def add(self, asset, path, copy=True):
 
         Args:
             asset (str): Name of asset.
-            path (str): Location on disk.
+            path (Path): Location on disk.
             copy (bool): Copy or move the content.
 
         Returns:
             str: Path to the asset on the filesystem.
         """
         assert isinstance(asset, str)
-        assert isinstance(path, str)
+        assert isinstance(path, Path)
         assert self.path, "cleanup() was called"
-        if not exists(path):
-            raise OSError(f"{path!r} does not exist")
-        path = abspath(path)
-        # only copy files from outside working path
-        if path.startswith(self.path):
-            raise AssetError(f"Cannot add existing asset content {path!r}")
-        dst_path = pathjoin(self.path, basename(path))
-        # avoid overwriting data that is part of an existing asset
-        if exists(dst_path):
-            raise AssetError(f"{basename(path)!r} is an existing asset")
+        if not path.exists():
+            raise OSError(f"'{path}' does not exist")
+        dst = self.path / path.name
         # remove existing asset with the same name
         if asset in self.assets:
             LOG.debug("asset %r exists, removing existing", asset)
             self.remove(asset)
+        # avoid overwriting data that is part of an existing asset
+        if dst.exists():
+            raise AssetError(f"{asset}: '{path.name}' already exists")
         if copy:
-            if isfile(path):
-                copyfile(path, dst_path)
+            if path.is_file():
+                copyfile(path, dst)
             else:
-                copytree(path, dst_path)
+                copytree(path, dst)
         else:
-            move(path, self.path)
-        self.assets[asset] = basename(path)
-        LOG.debug(
-            "added asset %r %s to %r", asset, "copied" if copy else "moved", dst_path
-        )
-        return dst_path
+            # TODO: move() only accepts str in Python 3.8
+            move(str(path), str(self.path))
+        self.assets[asset] = path.name
+        LOG.debug("%s asset %r to '%s'", "copied" if copy else "moved", asset, dst)
+        return dst
 
     def add_batch(self, assets):
         """Add collection of assets to the AssetManager.
@@ -84,7 +76,7 @@ def add_batch(self, assets):
             None
         """
         for asset, path in assets:
-            self.add(asset, path)
+            self.add(asset, Path(path))
 
     def cleanup(self):
         """Remove asset files from filesystem.
@@ -100,46 +92,17 @@ def cleanup(self):
             self.assets.clear()
             self.path = None
 
-    def dump(self, dst_path, subdir="_assets_"):
-        """Copy assets content to a given path.
-
-        Args:
-            dst_path (str): Path to copy assets content to.
-            subdir (str): Create and use as destination if given.
-
-        Returns:
-            dict: Collection asset paths keyed by asset name.
-        """
-        dumped = {}
-        if self.assets:
-            if subdir:
-                dst_path = pathjoin(dst_path, subdir)
-            makedirs(dst_path, exist_ok=not subdir)
-            for name, path in self.assets.items():
-                dumped[name] = path
-                src = pathjoin(self.path, path)
-                if isfile(src):
-                    copyfile(src, pathjoin(dst_path, path))
-                elif isdir(src):
-                    copytree(src, pathjoin(dst_path, path))
-                else:
-                    dumped.pop(name)
-                    LOG.warning("Failed to dump asset %r from %r", name, src)
-        return dumped
-
     def get(self, asset):
         """Get path to content on filesystem for given asset.
 
         Args:
             asset (str): Asset to lookup.
 
         Returns:
-            str: Path to asset content or None if asset does not exist.
+            Path: Path to asset content or None if asset does not exist.
         """
         item = self.assets.get(asset, None)
-        if item is not None:
-            return pathjoin(self.path, item)
-        return None
+        return self.path / item if item else None
 
     def is_empty(self):
         """Check if AssetManager contains entries.
@@ -159,15 +122,15 @@ def load(cls, assets, src_path, base_path=None):
         Args:
             asset (dict): Asset paths on filesystem relative to src_path, keyed on
                           asset name.
-            src_path (str): Path to scan for assets.
+            src_path (Path): Path to scan for assets.
             base_path (str): Base path to use to create local storage.
 
         Returns:
             AssetManager: Populated with contents provided by assets argument.
         """
         obj = cls(base_path=base_path)
         for asset, src_name in assets.items():
-            obj.add(asset, pathjoin(src_path, src_name))
+            obj.add(asset, src_path / src_name)
         return obj
 
     def remove(self, asset):
@@ -179,10 +142,10 @@ def remove(self, asset):
         Returns:
             None
         """
-        path = self.assets.pop(asset, None)
-        if path:
-            path = pathjoin(self.path, path)
-            if isfile(path):
-                unlink(path)
+        local_path = self.assets.pop(asset, None)
+        if local_path:
+            path = self.path / local_path
+            if path.is_file():
+                path.unlink()
             else:
                 rmtree(path, ignore_errors=True)

grizzly/target/puppet_target.py

@@ -383,11 +383,11 @@ def process_assets(self):
                 prefs = Path(tmp_path) / "prefs.js"
                 template = PrefPicker.lookup_template("browser-fuzzing.yml")
                 PrefPicker.load_template(template).create_prefsjs(prefs)
-                self._prefs = self.assets.add("prefs", str(prefs), copy=False)
+                self._prefs = self.assets.add("prefs", prefs, copy=False)
         abort_tokens = self.assets.get("abort-tokens")
         if abort_tokens:
             LOG.debug("loading 'abort tokens' from %r", abort_tokens)
-            with (Path(self.assets.path) / abort_tokens).open() as in_fp:
+            with (self.assets.path / abort_tokens).open() as in_fp:
                 for line in in_fp:
                     line = line.strip()
                     if line:
@@ -402,7 +402,7 @@ def process_assets(self):
             opts.load_options(self.environ.get(var_name, ""))
             if self.assets.get(asset):
                 # use suppression file if provided as asset
-                opts.add("suppressions", repr(self.assets.get(asset)), overwrite=True)
+                opts.add("suppressions", f"'{self.assets.get(asset)}'", overwrite=True)
             elif opts.get("suppressions"):
                 supp_file = opts.pop("suppressions")
                 if SanitizerOptions.is_quoted(supp_file):
@@ -412,7 +412,7 @@ def process_assets(self):
                     LOG.debug("using %r from environment", asset)
                     opts.add(
                         "suppressions",
-                        repr(self.assets.add(asset, supp_file)),
+                        f"'{self.assets.add(asset, Path(supp_file))}'",
                         overwrite=True,
                     )
                 else: