Skip to content

Commit

Permalink
Make FSSpecTarget be 'pathlike' (#474)
Browse files Browse the repository at this point in the history
The goal here is for recipe authors to be able to
pass *strings* to `target` parameters of StoreToZarr to specify
where the data should be for this specific recipe, and we can
inject the prefix with pangeo-forge-runner.

The primary use case is when we have a dictionary of recipes
being returned from a single recipe file. The *name* of the recipe
is unfortunately not really available during AST parsing, so the
suffix the data should be put in needs to be explicitly specified.
  • Loading branch information
yuvipanda authored Jan 17, 2023
1 parent 4b5f124 commit 241ca87
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
12 changes: 11 additions & 1 deletion pangeo_forge_recipes/storage.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

import hashlib
import io
import json
Expand All @@ -11,7 +13,7 @@
import unicodedata
from abc import ABC, abstractmethod
from contextlib import contextmanager
from dataclasses import dataclass
from dataclasses import dataclass, replace
from typing import Iterator, Optional, Sequence, Union
from urllib.parse import parse_qs, urlencode, urlparse, urlunparse

Expand Down Expand Up @@ -94,6 +96,14 @@ class FSSpecTarget(AbstractTarget):
fs: fsspec.AbstractFileSystem
root_path: str = ""

def __truediv__(self, suffix: str) -> FSSpecTarget:
"""
Support / operator so FSSpecTarget actslike a pathlib.path
Only supports getting a string suffix added in.
"""
return replace(self, root_path=os.path.join(self.root_path, suffix))

@classmethod
def from_url(cls, url: str):
fs, _, root_paths = fsspec.get_fs_token_paths(url)
Expand Down
6 changes: 6 additions & 0 deletions tests/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -98,3 +98,9 @@ def test_caching_only_truncates_long_fnames_for_local_fs(fs_cls, fname_longer_th
assert len(fname_in_full_path) == POSIX_MAX_FNAME_LENGTH
else:
assert len(fname_in_full_path) > POSIX_MAX_FNAME_LENGTH


def test_suffix(tmp_path):
assert str((FSSpecTarget(LocalFileSystem(), tmp_path) / "test").root_path) == str(
tmp_path / "test"
)

0 comments on commit 241ca87

Please sign in to comment.