Skip to content

Commit

Permalink
spec: wrap/override fsspec httpfs
Browse files Browse the repository at this point in the history
  • Loading branch information
pmrowla committed Feb 8, 2023
1 parent 3dbadcd commit 8582d20
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 23 deletions.
29 changes: 8 additions & 21 deletions dvc_http/__init__.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,8 @@
import threading
from getpass import getpass
from typing import TYPE_CHECKING, BinaryIO, Union
from typing import TYPE_CHECKING, Union

from dvc_objects.fs.base import AnyFSPath, FileSystem
from dvc_objects.fs.callbacks import DEFAULT_CALLBACK, Callback
from dvc_objects.fs.base import FileSystem
from dvc_objects.fs.errors import ConfigError
from funcy import cached_property, memoize, wrap_with

Expand Down Expand Up @@ -45,8 +44,8 @@ class HTTPFileSystem(FileSystem):

def __init__(self, fs=None, timeout=REQUEST_TIMEOUT, **kwargs):
super().__init__(fs, **kwargs)
self.upload_method = kwargs.get("method", "POST")

self.fs_args["upload_method"] = kwargs.get("method", "POST")
client_kwargs = self.fs_args.setdefault("client_kwargs", {})
client_kwargs.update(
{
Expand Down Expand Up @@ -129,28 +128,16 @@ async def get_client(self, **kwargs):

@cached_property
def fs(self):
from fsspec.implementations.http import (
HTTPFileSystem as _HTTPFileSystem,
)
from .spec import HTTPFileSystem as _HTTPFileSystem

return _HTTPFileSystem(get_client=self.get_client, **self.fs_args)
return _HTTPFileSystem(
get_client=self.get_client,
**self.fs_args,
)

def unstrip_protocol(self, path: str) -> str:
return path

def put_file(
self,
from_file: Union[AnyFSPath, BinaryIO],
to_info: AnyFSPath,
callback: Callback = DEFAULT_CALLBACK,
size: int = None,
**kwargs,
) -> None:
kwargs.setdefault("method", self.upload_method)
super().put_file(
from_file, to_info, callback=callback, size=size, **kwargs
)

# pylint: disable=arguments-differ

def find(self, *args, **kwargs):
Expand Down
11 changes: 11 additions & 0 deletions dvc_http/spec.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
from fsspec.implementations import http


class HTTPFileSystem(http.HTTPFileSystem): # pylint: disable=abstract-method
def __init__(self, *args, upload_method: str = "post", **kwargs):
super().__init__(*args, **kwargs)
self.upload_method = upload_method

async def _put_file(self, *args, **kwargs):
kwargs.setdefault("method", self.upload_method)
return await super()._put_file(*args, **kwargs)
6 changes: 4 additions & 2 deletions dvc_http/tests/test_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,10 +91,12 @@ def test_http_method():
}

fs = HTTPFileSystem(**config, method="PUT")
assert fs.upload_method == "PUT"
assert fs.fs_args.get("upload_method") == "PUT"
assert fs.fs.upload_method == "PUT"

fs = HTTPFileSystem(**config, method="POST")
assert fs.upload_method == "POST"
assert fs.fs_args.get("upload_method") == "POST"
assert fs.fs.upload_method == "POST"


@pytest.mark.parametrize(
Expand Down

0 comments on commit 8582d20

Please sign in to comment.