diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml index 39e5987..0465129 100644 --- a/.github/workflows/run_tests.yml +++ b/.github/workflows/run_tests.yml @@ -42,6 +42,11 @@ jobs: # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide flake8 . --count --exit-zero --ignore=E402,C901 --max-line-length=127 --statistics + - name: Check types with mypy + continue-on-error: true + run: | + python -m mypy --install-types --non-interactive cottoncandy/ + - name: Test with pytest env: DL_BUCKET_NAME: ${{ secrets.DL_BUCKET_NAME }} diff --git a/cottoncandy/__init__.py b/cottoncandy/__init__.py index fd62b79..dcee2a7 100644 --- a/cottoncandy/__init__.py +++ b/cottoncandy/__init__.py @@ -2,33 +2,34 @@ ''' +from typing import Literal import os from cottoncandy import options +from .browser import BrowserObject +from .interfaces import DefaultInterface from .utils import get_keys, string2bool __version__ = "0.4.0" -ACCESS_KEY = options.config.get('login', 'access_key') -SECRET_KEY = options.config.get('login', 'secret_key') -ENDPOINT_URL = options.config.get('login', 'endpoint_url') +ACCESS_KEY: str = options.config.get('login', 'access_key') +SECRET_KEY: str = options.config.get('login', 'secret_key') +ENDPOINT_URL: str = options.config.get('login', 'endpoint_url') DEFAULT_SIGNATURE_VERSION = options.config.get('basic', 'signature_version') -default_bucket = options.config.get('basic', 'default_bucket') -force_bucket_creation = options.config.get('basic', 'force_bucket_creation') -force_bucket_creation = string2bool(force_bucket_creation) - - -def get_interface(bucket_name=default_bucket, - ACCESS_KEY=ACCESS_KEY, - SECRET_KEY=SECRET_KEY, - endpoint_url=ENDPOINT_URL, - force_bucket_creation=force_bucket_creation, - verbose=True, - backend='s3', - **kwargs): +default_bucket: str = options.config.get('basic', 'default_bucket') +force_bucket_creation: bool = bool(string2bool(options.config.get('basic', 'force_bucket_creation'))) + +def get_interface(bucket_name: str=default_bucket, + ACCESS_KEY: str=ACCESS_KEY, + SECRET_KEY: str=SECRET_KEY, + endpoint_url: str=ENDPOINT_URL, + force_bucket_creation: bool=force_bucket_creation, + verbose: bool=True, + backend: Literal['s3', 'gdrive', 'local']='s3', + **kwargs) -> DefaultInterface: """Return an interface to the cloud. Parameters @@ -38,7 +39,7 @@ def get_interface(bucket_name=default_bucket, SECRET_KEY : str endpoint_url : str The URL for the S3 gateway - backend : 's3'|'gdrive' + backend : 's3'|'gdrive'|'local' What backend to hook on to kwargs : S3 only. kwargs passed to botocore. For example, @@ -48,7 +49,7 @@ def get_interface(bucket_name=default_bucket, Returns ------- - cci : cottoncandy.InterfaceObject + cci : cottoncandy.DefaultInterface """ from cottoncandy.interfaces import DefaultInterface @@ -82,10 +83,10 @@ def get_interface(bucket_name=default_bucket, return interface -def get_browser(bucket_name=default_bucket, - ACCESS_KEY=ACCESS_KEY, - SECRET_KEY=SECRET_KEY, - endpoint_url=ENDPOINT_URL): +def get_browser(bucket_name: str=default_bucket, + ACCESS_KEY: str=ACCESS_KEY, + SECRET_KEY: str=SECRET_KEY, + endpoint_url: str=ENDPOINT_URL) -> BrowserObject: """Browser object that allows you to tab-complete your way through your objects @@ -131,4 +132,4 @@ def get_browser(bucket_name=default_bucket, return S3Directory('/', interface=interface) -__all__ = ['get_interface', 'get_browser', 'interfaces', 'browser'] +__all__ = ['get_interface', 'get_browser', 'interfaces', 'browser', 'DefaultInterface'] diff --git a/cottoncandy/backend.py b/cottoncandy/backend.py index 1a49831..7de1ecf 100644 --- a/cottoncandy/backend.py +++ b/cottoncandy/backend.py @@ -1,10 +1,23 @@ from abc import ABCMeta, abstractmethod +from typing import Any, NamedTuple, BinaryIO, Optional, Sequence class FileNotFoundError(RuntimeError): """File not found error""" +class CloudStream(NamedTuple): + """ + A simple unified representation of an object downloaded from the cloud. + .content is a streaming object with a .read() function + .metadata is a dictionary of the custom metadata of this object + + TODO: unified metadata + """ + content: BinaryIO + metadata: dict[str, str] + + class CCBackEnd: """ Interface for cottoncandy backends @@ -14,14 +27,20 @@ class CCBackEnd: def __init__(self): pass + @property + @abstractmethod + def bucket_name(self) -> Optional[str]: + """Name of the current bucket or backing path.""" + pass + ## Basic File IO @abstractmethod - def check_file_exists(self, file_name, bucket_name): + def check_file_exists(self, cloud_name: str, bucket_name: Optional[str] = None) -> bool: """Checks whether a file exists on the cloud Parameters ---------- - file_name : str + cloud_name : str file on cloud bucket_name : str (s3) bucket to check in @@ -33,7 +52,7 @@ def check_file_exists(self, file_name, bucket_name): pass @abstractmethod - def upload_stream(self, stream, cloud_name, metadata, permissions, threads): + def upload_stream(self, stream: BinaryIO, cloud_name: str, metadata: dict[str, str], permissions: Optional[str], threads: int) -> None: """Uploads a stream object with a .read() function Parameters @@ -52,12 +71,12 @@ def upload_stream(self, stream, cloud_name, metadata, permissions, threads): Returns ------- - bool, upload success + None """ pass @abstractmethod - def upload_file(self, file_name, cloud_name, permissions, threads): + def upload_file(self, file_name: str, cloud_name: Optional[str] = None, permissions: Optional[str] = None, threads: int = 1) -> None: """Uploads a file from disk Parameters @@ -74,13 +93,13 @@ def upload_file(self, file_name, cloud_name, permissions, threads): Returns ------- - bool, upload success + None """ pass @abstractmethod - def download_stream(self, cloud_name, threads): + def download_stream(self, cloud_name: str, threads: int) -> CloudStream: """Downloads a object to an in-memory stream Parameters @@ -98,7 +117,7 @@ def download_stream(self, cloud_name, threads): pass @abstractmethod - def download_to_file(self, cloud_name, file_name, threads): + def download_to_file(self, cloud_name: str, file_name: str, threads: int) -> None: """Downloads an object directly to disk Parameters @@ -113,14 +132,14 @@ def download_to_file(self, cloud_name, file_name, threads): Returns ------- - bool, download success + None """ pass ## Basic File management @abstractmethod - def list_directory(self, path, limit): + def list_directory(self, path: str, limit: int) -> list[str]: """Lists the content of a directory Parameters @@ -130,12 +149,12 @@ def list_directory(self, path, limit): Returns ------- - + list[str] """ pass @abstractmethod - def list_objects(self): + def list_objects(self) -> Sequence[Any]: """Gets all objects contained by backend Returns @@ -145,7 +164,7 @@ def list_objects(self): pass @abstractmethod - def copy(self, source, destination, source_bucket, destination_bucket, overwrite): + def copy(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False) -> Any: """Copies an object Parameters @@ -168,25 +187,30 @@ def copy(self, source, destination, source_bucket, destination_bucket, overwrite pass @abstractmethod - def move(self, source, destination, source_bucket, destination_bucket, overwrite): + def move(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False) -> Any: """Moves an object Parameters ---------- - source - destination - source_bucket - destination_bucket - overwrite + source: str + origin path + destination: str + destination path + source_bucket: str + (s3) origin bucket + destination_bucket: str + (s3) destination bucket + overwrite: bool + overwrite if destination exists? Returns ------- - + bool, move success """ pass @abstractmethod - def delete(self, file_name, recursive=False, delete=False): + def delete(self, cloud_name: str, recursive: bool = False, delete: bool = False) -> bool: """Deletes an object Parameters @@ -200,13 +224,13 @@ def delete(self, file_name, recursive=False, delete=False): Returns ------- - + bool, delete success """ pass @property @abstractmethod - def size(self): + def size(self) -> int: """Size of stored cloud items in bytes Returns @@ -214,16 +238,3 @@ def size(self): int """ pass - - -class CloudStream: - """ - A simple unified representation of an object downloaded from the cloud. - .content is a streaming object with a .read() function - .metadata is a dictionary of the custom metadata of this object - - TODO: unified metadata - """ - def __init__(self, stream, metadata): - self.content = stream - self.metadata = metadata diff --git a/cottoncandy/browser.py b/cottoncandy/browser.py index 82dc43f..7924b70 100644 --- a/cottoncandy/browser.py +++ b/cottoncandy/browser.py @@ -165,7 +165,7 @@ def __repr__(self): def __len__(self): return len(self._subdirs) - def __dir__(self): + def __dir__(self): # type: ignore[no-redef] return list(self._subdirs.keys()) def __getattr__(self, attr): diff --git a/cottoncandy/gdriveclient.py b/cottoncandy/gdriveclient.py index dd11a2b..a3a570c 100644 --- a/cottoncandy/gdriveclient.py +++ b/cottoncandy/gdriveclient.py @@ -15,11 +15,11 @@ except ImportError: try: # support >=ipython-0.11, int: files = self.drive.ListFile({'q': "trashed=false"}).GetList() sizes = [f.metadata['size'] for f in files] return sum(sizes) diff --git a/cottoncandy/interfaces.py b/cottoncandy/interfaces.py index 8fc54b8..3ca8d9c 100644 --- a/cottoncandy/interfaces.py +++ b/cottoncandy/interfaces.py @@ -9,9 +9,10 @@ from warnings import warn import six +from typing import Any, BinaryIO, Iterable, Literal, Mapping, Sequence, TypedDict, Optional, Union, cast, no_type_check import cottoncandy.browser -from cottoncandy.backend import FileNotFoundError +from cottoncandy.backend import FileNotFoundError, CloudStream from .options import config from .s3client import S3Client, botocore @@ -39,20 +40,23 @@ ) DO_COMPRESSION = config.get('compression', 'do_compression').lower() in ('true', 't', 'y', 'yes') -COMPRESSION_SMALL = config.get('compression', 'small_array') -COMPRESSION_LARGE = config.get('compression', 'large_array') +COMPRESSION_SMALL: str = config.get('compression', 'small_array') +COMPRESSION_LARGE: str = config.get('compression', 'large_array') +import numpy as np +import numpy.typing as npt try: - import numpy as np - from scipy.sparse import bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix + import scipy.sparse except ImportError: - warn('numpy/scipy not available') + warn('scipy not available') try: import numcodecs except ImportError: warn('numcodecs python library not available') +NestedArrayDict = dict[str, Union[npt.NDArray, None, 'NestedArrayDict']] + # ------------------ # Cloud Interfaces @@ -66,10 +70,10 @@ class BasicInterface(InterfaceObject): """Basic cottoncandy interface to the cloud. """ - def __init__(self, bucket_name, - ACCESS_KEY, SECRET_KEY, url=None, - force_bucket_creation=False, - verbose=True, backend='s3', **kwargs): + def __init__(self, bucket_name: Union[str, None], + ACCESS_KEY: str, SECRET_KEY: str, url: str, + force_bucket_creation: bool = False, + verbose: bool = True, backend: str='s3', **kwargs): """ Parameters ---------- @@ -102,10 +106,11 @@ def __init__(self, bucket_name, **kwargs) elif backend == 'gdrive': from .gdriveclient import GDriveClient - self.backend_interface = GDriveClient(ACCESS_KEY, SECRET_KEY) + self.backend_interface = GDriveClient(ACCESS_KEY, SECRET_KEY) # type: ignore[assignment] elif backend == 'local': from .localclient import LocalClient - self.backend_interface = LocalClient(path=bucket_name) + assert bucket_name is not None, "Must specify bucket_name for 'local' backend (this is the local path to use)" + self.backend_interface = LocalClient(path=bucket_name) # type: ignore[assignment] else: raise ValueError('Bad backend') @@ -134,21 +139,15 @@ def __repr__(self): def _get_bucket_name(self, bucket_name): return self.backend_interface._get_bucket_name(bucket_name) - def pathjoin(self, a, *p): + def pathjoin(self, a: str, *p: str) -> str: return pathjoin(a, *p) @property - def bucket_name(self): - if self.backend == "s3": - return self.backend_interface.bucket_name - elif self.backend == "gdrive": - print('Google drive has no concept of buckets') - return None - else: - return self.backend_interface.path + def bucket_name(self) -> Optional[str]: + return self.backend_interface.bucket_name @clean_object_name - def exists_object(self, object_name, bucket_name=None, raise_err=False): + def exists_object(self, object_name: str, bucket_name: Optional[str]=None, raise_err: bool=False) -> bool: """Check whether object exists in bucket Parameters @@ -165,15 +164,15 @@ def exists_object(self, object_name, bucket_name=None, raise_err=False): else: return exists - def exists_bucket(self, bucket_name): + def exists_bucket(self, bucket_name: str) -> bool: """Check whether the bucket exists""" return self.backend_interface.check_bucket_exists(bucket_name) - def create_bucket(self, bucket_name, acl=DEFAULT_ACL): + def create_bucket(self, bucket_name: str, acl=DEFAULT_ACL) -> None: """Create a new bucket""" self.backend_interface.create_bucket(bucket_name, acl) - def rm_bucket(self, bucket_name): + def rm_bucket(self, bucket_name: str) -> None: '''Remove an empty bucket. Throws an exception when bucket is not empty. ''' self.set_bucket(bucket_name) @@ -183,7 +182,7 @@ def rm_bucket(self, bucket_name): except botocore.exceptions.ClientError: print("Bucket not empty. To delete, first empty the bucket.") - def set_bucket(self, bucket_name): + def set_bucket(self, bucket_name: str) -> None: """Bucket to use""" self.backend_interface.set_current_bucket(bucket_name) @@ -191,7 +190,7 @@ def get_bucket(self): """Get bucket boto3 object""" return self.backend_interface.get_bucket() - def get_bucket_objects(self, **kwargs): + def get_bucket_objects(self, **kwargs) -> Sequence[Any]: """Get list of objects from the bucket. This is a wrapper to ``self.get_bucket().bucket.objects`` @@ -221,7 +220,7 @@ def get_bucket_objects(self, **kwargs): warn('Deprecated. Use get_objects() instead', DeprecationWarning) return self.backend_interface.list_objects(**kwargs) - def get_objects(self, **kwargs): + def get_objects(self, **kwargs) -> Sequence[Any]: """ Like get_bucket_objects, but more aptly named to the generic interface Parameters @@ -235,7 +234,7 @@ def get_objects(self, **kwargs): """ return self.backend_interface.list_objects(**kwargs) - def get_bucket_size(self, limit=10**6, page_size=10**6): + def get_bucket_size(self, limit: int=10**6, page_size: int=10**6) -> int: """Counts the size of all objects in the current bucket. Parameters @@ -261,7 +260,7 @@ def get_bucket_size(self, limit=10**6, page_size=10**6): warn('Deprecated, use get_size() instead', DeprecationWarning) return self.backend_interface.size - def get_size(self): + def get_size(self) -> int: """ Gets the total size of the current container of objects. Generic naming. Parameters @@ -274,17 +273,17 @@ def get_size(self): """ return self.backend_interface.size - def show_buckets(self): + def show_buckets(self) -> None: """Show available buckets""" self.backend_interface.show_all_buckets() @clean_object_name - def get_object(self, object_name, bucket_name=None): + def get_object(self, object_name: str, bucket_name: Optional[str]=None): """Get a boto3 object. Create it if it doesn't exist""" # NOTE: keeping this in case outside code is using this. return self.backend_interface.get_s3_object(object_name, bucket_name) - def show_objects(self, limit=1000, page_size=1000): + def show_objects(self, limit: int=1000, page_size: int=1000) -> None: """Print objects in the current bucket""" if self.backend == 's3': object_list = self.backend_interface.list_objects(limit=limit, page_size=page_size) @@ -295,7 +294,7 @@ def show_objects(self, limit=1000, page_size=1000): object_list = self.backend_interface.list_objects(limit = limit, page_size = page_size * 100) print_objects(object_list) elif self.backend == 'gdrive': - drivefiles = self.backend_interface.drive.ListFile({'q': "trashed=false"}).GetList() + drivefiles = self.backend_interface.drive.ListFile({'q': "trashed=false"}).GetList() # type: ignore object_list = [df['title'] for df in drivefiles] for obj in object_list: # TODO: also print last modified date and whatever else to match s3 @@ -308,12 +307,11 @@ def show_objects(self, limit=1000, page_size=1000): # CCBackEnd object or the actual cloud APIs @clean_object_name - def upload_object(self, object_name, body, acl=DEFAULT_ACL, threads = THREADS, **metadata): + def upload_object(self, object_name: str, body: BinaryIO, acl: str=DEFAULT_ACL, threads: int=THREADS, **metadata: str) -> None: # First check size of object to see if MPU is necessary - self.backend_interface.upload_stream(body, object_name, metadata, permissions = acl, threads = threads) - def download_stream(self, object_name, threads = THREADS): + def download_stream(self, object_name: str, threads: int = THREADS) -> CloudStream: """ Returns the CloudStream object for an object Parameters @@ -329,9 +327,9 @@ def download_stream(self, object_name, threads = THREADS): """ return self.backend_interface.download_stream(object_name, threads) - def upload_from_file(self, flname, object_name=None, - ExtraArgs=dict(ACL=DEFAULT_ACL), - threads = THREADS): + def upload_from_file(self, flname: str, object_name: Optional[str]=None, + ExtraArgs: Optional[Mapping[str, str]]=None, + threads: int = THREADS) -> None: """Upload a file to the cloud. Parameters @@ -350,10 +348,11 @@ def upload_from_file(self, flname, object_name=None, ------- response : boto3 response """ - return self.backend_interface.upload_file(flname, object_name, ExtraArgs['ACL'], threads) + extra_args = dict(ACL=DEFAULT_ACL) if ExtraArgs is None else ExtraArgs + return self.backend_interface.upload_file(flname, object_name, extra_args['ACL'], threads) - def upload_from_directory(self, disk_path, cloud_path=None, - recursive=False, ExtraArgs=dict(ACL=DEFAULT_ACL), threads = THREADS): + def upload_from_directory(self, disk_path: str, cloud_path: Optional[str]=None, + recursive: bool=False, ExtraArgs=dict(ACL=DEFAULT_ACL), threads: int = THREADS) -> None: '''Upload a directory to the cloud ''' filenames = sorted(os.listdir(disk_path)) @@ -373,7 +372,7 @@ def upload_from_directory(self, disk_path, cloud_path=None, print('Uploaded "%s" to "%s"' % (disk_path, cloud_path)) @clean_object_name - def download_to_file(self, object_name, file_name, threads = THREADS): + def download_to_file(self, object_name: str, file_name: str, threads: int = THREADS) -> None: """Download cloud object to a file Parameters @@ -387,7 +386,7 @@ def download_to_file(self, object_name, file_name, threads = THREADS): return self.backend_interface.download_to_file(object_name, file_name, threads) @clean_object_name - def download_object(self, object_name, threads = THREADS): + def download_object(self, object_name: str, threads: int = THREADS) -> Any: """Download object raw data. This simply calls the object body ``read()`` method. @@ -405,7 +404,7 @@ def download_object(self, object_name, threads = THREADS): return self.download_stream(object_name, threads).content.read() @clean_object_name - def upload_json(self, object_name, ddict, acl=DEFAULT_ACL, threads = 1, **metadata): + def upload_json(self, object_name: str, ddict: Mapping[Any, Any], acl: str = DEFAULT_ACL, threads: int = 1, **metadata: str) -> None: """Upload a dict as a JSON using ``json.dumps`` Parameters @@ -419,7 +418,7 @@ def upload_json(self, object_name, ddict, acl=DEFAULT_ACL, threads = 1, **metada return self.upload_object(object_name, StringIO(json_data.encode()), acl, threads, **metadata) @clean_object_name - def download_json(self, object_name, threads = 1): + def download_json(self, object_name: str, threads: int = 1) -> Any: """Download a JSON object Parameters @@ -437,7 +436,7 @@ def download_json(self, object_name, threads = 1): return json.loads(obj.decode()) @clean_object_name - def upload_pickle(self, object_name, data_object, acl=DEFAULT_ACL, threads = THREADS, **metadata): + def upload_pickle(self, object_name: str, data_object: Any, acl: str=DEFAULT_ACL, threads: int = THREADS, **metadata) -> None: """Upload an object using pickle: ``pickle.dumps`` Parameters @@ -452,7 +451,7 @@ def upload_pickle(self, object_name, data_object, acl=DEFAULT_ACL, threads = THR return response @clean_object_name - def download_pickle(self, object_name, threads = THREADS): + def download_pickle(self, object_name: str, threads: int = THREADS) -> Any: """Download a pickle object Parameters @@ -494,7 +493,7 @@ def __init__(self, *args, **kwargs): super(ArrayInterface, self).__init__(*args, **kwargs) @clean_object_name - def upload_npy_array(self, object_name, array, acl=DEFAULT_ACL, threads = THREADS, **metadata): + def upload_npy_array(self, object_name: str, array: npt.NDArray, acl: str=DEFAULT_ACL, threads: int = THREADS, **metadata: str) -> None: """Upload a np.ndarray using ``np.save`` This method creates a copy of the array in memory @@ -526,7 +525,7 @@ def upload_npy_array(self, object_name, array, acl=DEFAULT_ACL, threads = THREAD return response @clean_object_name - def download_npy_array(self, object_name, threads = THREADS): + def download_npy_array(self, object_name: str, threads: int = THREADS) -> npt.NDArray: """Download a np.ndarray uploaded using ``np.save`` with ``np.load``. Parameters @@ -544,7 +543,7 @@ def download_npy_array(self, object_name, threads = THREADS): return array @clean_object_name - def upload_raw_array(self, object_name, array, compression=DO_COMPRESSION, acl=DEFAULT_ACL, threads = THREADS, **metadata): + def upload_raw_array(self, object_name: str, array: npt.NDArray, compression: Optional[Union[bool, str]]=DO_COMPRESSION, acl: str=DEFAULT_ACL, threads: int = THREADS, **metadata: str) -> None: """Upload a binary representation of a np.ndarray This method reads the array content from memory to upload. @@ -575,8 +574,8 @@ def upload_raw_array(self, object_name, array, compression=DO_COMPRESSION, acl=D # Backward compatibility if 'gzip' in metadata: warn("Deprecated keyword argument `gzip`. Use `compression='gzip'` instead", DeprecationWarning) - gz = metadata.pop('gzip') - compression = 'gzip' if gz else False + use_gzip = metadata.pop('gzip') + compression = 'gzip' if use_gzip else False if compression is True: # check whether array is >= 2 GB @@ -588,8 +587,8 @@ def upload_raw_array(self, object_name, array, compression=DO_COMPRESSION, acl=D raise ValueError("gzip does not support compression of >2GB arrays. " "Try `compression='Zstd'` instead.") - order = 'C' if array.flags.carray else 'F' - if ((not array.flags['%s_CONTIGUOUS' % order] and six.PY2) or + order: Literal['C', 'F'] = 'C' if array.flags.carray else 'F' + if ((not array.flags['%s_CONTIGUOUS' % order] and six.PY2) or # type: ignore (not array.flags['C_CONTIGUOUS'] and six.PY3)): warn('Non-contiguous array. Creating copy (will use extra memory)...') @@ -642,7 +641,7 @@ def upload_raw_array(self, object_name, array, compression=DO_COMPRESSION, acl=D return response @clean_object_name - def download_raw_array(self, object_name, buffersize=2**16, threads = THREADS, **kwargs): + def download_raw_array(self, object_name: str, buffersize: int=2**16, threads: int = THREADS, **kwargs) -> npt.NDArray: """Download a binary np.ndarray and return an np.ndarray object This method downloads an array without any disk or memory overhead. @@ -669,7 +668,8 @@ def download_raw_array(self, object_name, buffersize=2**16, threads = THREADS, * shape = arraystream.metadata['shape'] shape = tuple(map(int, shape.split(',')) if shape else ()) dtype = np.dtype(arraystream.metadata['dtype']) - order = arraystream.metadata.get('order', 'C') + order = cast(Literal['C', 'F'], arraystream.metadata.get('order', 'C')) + assert order in ['C', 'F'], f'Invalid array order in metadata: {order}' array = np.empty(tuple(shape), dtype = dtype, order = order) body = arraystream.content @@ -685,6 +685,7 @@ def download_raw_array(self, object_name, buffersize=2**16, threads = THREADS, * assert 'compression' in arraystream.metadata + datastream: BinaryIO if arraystream.metadata['compression'] == 'gzip': # gzipped! datastream = GzipInputStream(body) @@ -706,14 +707,14 @@ def download_raw_array(self, object_name, buffersize=2**16, threads = THREADS, * return array @clean_object_name - def dict2cloud(self, object_name, array_dict, acl=DEFAULT_ACL, - verbose=True, threads = THREADS, **metadata): + def dict2cloud(self, object_name: str, array_dict: NestedArrayDict, acl: str = DEFAULT_ACL, + verbose: bool = True, threads: int = THREADS, **metadata: str): """Upload an arbitrary depth dictionary containing arrays Parameters ---------- object_name : str - array_dict : dict + array_dict : dict[str, Union[npt.NDArray, 'NestedArrayDict']] An arbitrary depth dictionary of arrays. This can be conceptualized as implementing an HDF-like group verbose : bool @@ -725,17 +726,17 @@ def dict2cloud(self, object_name, array_dict, acl=DEFAULT_ACL, name = self.pathjoin(object_name, k) if isinstance(v, dict): - _ = self.dict2cloud(name, v, acl=acl, threads = threads, **metadata) + _ = self.dict2cloud(name, v, acl=acl, verbose=verbose, threads = threads, **metadata) elif isinstance(v, np.ndarray): _ = self.upload_raw_array(name, v, acl=acl, threads = threads, **metadata) else: # try converting to array - _ = self.upload_raw_array(name, np.asarray(v), threads = threads, acl=acl) + _ = self.upload_raw_array(name, np.asarray(v), threads = threads, acl=acl, **metadata) if verbose: print('uploaded arrays in "%s"' % object_name) @clean_object_name - def cloud2dict(self, object_root, verbose=True, keys=None, threads = THREADS, **metadata): + def cloud2dict(self, object_root: str, verbose: bool = True, keys=None, threads: int = THREADS, **metadata: str) -> NestedArrayDict: """Download all the arrays of the object branch and return a dictionary. This is the complement to ``dict2cloud`` @@ -756,20 +757,21 @@ def cloud2dict(self, object_root, verbose=True, keys=None, threads = THREADS, ** An arbitrary depth dictionary. """ # TODO: gdrive compatibility? - datadict = {} + datadict: NestedArrayDict = {} + subdirs: list[str] if keys is not None: if isinstance(keys, str): keys = [keys] subdirs = keys else: - subdirs = self.lsdir(object_root) + subdirs = self.lsdir(object_root) # type: ignore subdirs = [os.path.split(t)[-1] for t in subdirs] if not subdirs: print('Nothing found in "%s"' % object_root) - return + return datadict for subdir in subdirs: path = self.pathjoin(object_root, subdir) @@ -783,7 +785,7 @@ def cloud2dict(self, object_root, verbose=True, keys=None, threads = THREADS, ** arr = None datadict[subdir] = arr else: - datadict[subdir] = self.cloud2dict(path, threads = threads) + datadict[subdir] = self.cloud2dict(path, threads = threads, verbose = verbose, **metadata) if verbose: print('Downloaded arrays in "%s"' % object_root) @@ -791,7 +793,7 @@ def cloud2dict(self, object_root, verbose=True, keys=None, threads = THREADS, ** return datadict @clean_object_name - def cloud2dataset(self, object_root, **metadata): + def cloud2dataset(self, object_root: str, **metadata): """Get a dataset representation of the object branch. Parameters @@ -809,7 +811,7 @@ def cloud2dataset(self, object_root, **metadata): return S3Directory(object_root, interface = self) @clean_object_name - def upload_dask_array(self, object_name, arr, axis=-1, buffersize=DASK_CHUNKSIZE, threads = THREADS, **metakwargs): + def upload_dask_array(self, object_name: str, arr: npt.NDArray, axis: int = -1, buffersize: int = DASK_CHUNKSIZE, threads: int = THREADS, **metakwargs: str) -> None: """Upload an array in chunks and store the metadata to reconstruct the complete matrix with ``dask``. @@ -841,11 +843,20 @@ def upload_dask_array(self, object_name, arr, axis=-1, buffersize=DASK_CHUNKSIZE * my_array_name/pt0001 * my_array_name/metadata.json """ - metadata = dict(shape = arr.shape, - dtype = arr.dtype.str, - dask = [], - chunk_sizes = [], - ) + class DaskArrayMetadata(TypedDict): + shape: tuple[int, ...] + dtype: str + dask: list[tuple[tuple[int, ...], str]] + chunk_sizes: list[tuple[int, ...]] + chunks: list[list[int]] + + metadata: DaskArrayMetadata = { + 'shape': arr.shape, + 'dtype': arr.dtype.str, + 'dask': [], + 'chunk_sizes': [], + 'chunks': [], + } generator = generate_ndarray_chunks(arr, axis = axis, buffersize = buffersize) total_upload = 0.0 @@ -862,7 +873,7 @@ def upload_dask_array(self, object_name, arr, axis=-1, buffersize=DASK_CHUNKSIZE # convert to dask convention (sorry) details = [t[0] for t in metadata['dask']] - dimension_sizes = [dict() for idx in range(arr.ndim)] + dimension_sizes: list[dict[int, int]] = [dict() for idx in range(arr.ndim)] for dim, chunks in enumerate(zip(*details)): for sample_idx, chunk_idx in enumerate(chunks): if chunk_idx not in dimension_sizes[dim]: @@ -870,10 +881,10 @@ def upload_dask_array(self, object_name, arr, axis=-1, buffersize=DASK_CHUNKSIZE chunks = [[value for k, value in sorted(sizes.items())] for sizes in dimension_sizes] metadata['chunks'] = chunks - return self.upload_json(self.pathjoin(object_name, 'metadata.json'), metadata, **metakwargs) + return self.upload_json(self.pathjoin(object_name, 'metadata.json'), metadata, threads=threads, **metakwargs) @clean_object_name - def download_dask_array(self, object_name, dask_name='array', threads = THREADS): + def download_dask_array(self, object_name: str, dask_name: str = 'array', threads: int = THREADS) -> Any: """Downloads a split matrix as a ``dask.array.Array`` object This uses the stored object metadata to reconstruct the full @@ -905,7 +916,7 @@ def download_dask_array(self, object_name, dask_name='array', threads = THREADS) return da.Array(dask, dask_name, chunks, shape = shape, dtype = dtype) @clean_object_name - def upload_sparse_array(self, object_name, arr, threads = THREADS): + def upload_sparse_array(self, object_name: str, arr: Any, threads: int = THREADS) -> None: """Uploads a scipy.sparse array as a folder of array objects Parameters @@ -916,8 +927,11 @@ def upload_sparse_array(self, object_name, arr, threads = THREADS): A scipy.sparse array to be saved. If type is DOK or LIL, it will be converted to csr before saving threads: int - number of connection threads to use + number of connection threads to use """ + # Import sparse matrix classes here to avoid unbound reference + from scipy.sparse import bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix + if isinstance(arr, csr_matrix): attrs = ['data', 'indices', 'indptr'] arrtype = 'csr' @@ -948,7 +962,7 @@ def upload_sparse_array(self, object_name, arr, threads = THREADS): return self.upload_json(self.pathjoin(object_name, 'metadata.json'), metadata) @clean_object_name - def download_sparse_array(self, object_name, threads = THREADS): + def download_sparse_array(self, object_name: str, threads: int = THREADS) -> Any: """Downloads a scipy.sparse array Parameters @@ -956,13 +970,16 @@ def download_sparse_array(self, object_name, threads = THREADS): object_name : str The object name for the sparse array to be retrieved. threads: int - number of connection threads to use + number of connection threads to use Returns ------- arr : scipy.sparse.spmatrix The array stored at the location given by object_name """ + # Import sparse matrix classes here to avoid unbound reference + from scipy.sparse import bsr_matrix, coo_matrix, csc_matrix, csr_matrix, dia_matrix + # Get metadata metadata = self.download_json(self.pathjoin(object_name, 'metadata.json')) # Get type, shape @@ -973,6 +990,7 @@ def download_sparse_array(self, object_name, threads = THREADS): for attr in metadata['attrs']: d[attr] = self.download_raw_array(self.pathjoin(object_name, attr), threads = threads) + arr: Union[csr_matrix, coo_matrix, csc_matrix, bsr_matrix, dia_matrix] if arrtype == 'csr': arr = csr_matrix((d['data'], d['indices'], d['indptr']), shape = shape) @@ -987,6 +1005,8 @@ def download_sparse_array(self, object_name, threads = THREADS): shape = shape) elif arrtype == 'dia': arr = dia_matrix((d['data'], d['offsets']), shape = shape) + else: + raise ValueError(f"unsupported sparse array type: {arrtype}") return arr @@ -1014,7 +1034,7 @@ def __init__(self, *args, **kwargs): """ super(FileSystemInterface, self).__init__(*args, **kwargs) - def lsdir(self, path='/', limit=10**3): + def lsdir(self, path: str='/', limit: int=10**3) -> list[str]: """List the contents of a directory Parameters @@ -1029,7 +1049,7 @@ def lsdir(self, path='/', limit=10**3): return self.backend_interface.list_directory(path, limit) @clean_object_name - def ls(self, pattern, page_size=10**3, limit=10**3, verbose=False): + def ls(self, pattern: str, page_size: int=10**3, limit: int=10**3, verbose: bool=False) -> list[str]: """File-system like search for S3 objects Parameters @@ -1061,7 +1081,7 @@ def ls(self, pattern, page_size=10**3, limit=10**3, verbose=False): # get objects that match common prefix if not has_real_magic(pattern): - object_names = self.lsdir(prefix, limit = limit) + object_names: Iterable[str] = self.lsdir(prefix, limit = limit) else: object_list = self.get_objects(filter = dict(Prefix = prefix), page_size = page_size, @@ -1082,7 +1102,7 @@ def ls(self, pattern, page_size=10**3, limit=10**3, verbose=False): return list(object_names) @clean_object_name - def glob(self, pattern, **kwargs): + def glob(self, pattern: str, **kwargs) -> list[str]: """Return a list of object names in the cloud storage that match the glob pattern. @@ -1138,7 +1158,8 @@ def glob(self, pattern, **kwargs): else: return self.glob_s3(pattern, **kwargs) - def glob_google_drive(self, pattern): + @no_type_check # this function isn't implemented yet + def glob_google_drive(self, pattern: str) -> list[str]: """Globbing on google drive Parameters @@ -1154,7 +1175,10 @@ def glob_google_drive(self, pattern): return self.backend_interface.list_objects(True) nextFolderExp = r'^/?[^/]*/' - nextFolder = re.match(nextFolderExp, pattern).group(0) + match = re.match(nextFolderExp, pattern) + if match is None: + raise ValueError('Invalid pattern: %s' % pattern) + nextFolder = match.group(0) pattern = re.sub(nextFolderExp, '', pattern) if '*' not in nextFolder: # no wildcard, simply cd into it @@ -1165,16 +1189,17 @@ def glob_google_drive(self, pattern): matches.append(self.glob_google_drive()) # TODO: finish this + raise NotImplementedError('Globbing on google drive not yet implemented') - return + return matches - def glob_s3(self, pattern, **kwargs): + def glob_s3(self, pattern: str, **kwargs) -> list[str]: """Globbing on S3 Parameters ---------- - pattern + pattern: str kwargs Returns @@ -1212,7 +1237,7 @@ def glob_s3(self, pattern, **kwargs): return matches @clean_object_name - def download_directory(self, directory, disk_name): + def download_directory(self, directory: str, disk_name: os.PathLike) -> None: """ Download an entire directory NOTE: currently only tested on s3 @@ -1222,7 +1247,7 @@ def download_directory(self, directory, disk_name): self directory : str directory on s3 to download - disk_name : + disk_name : PathLike name of directory on disk to download to Returns @@ -1246,20 +1271,20 @@ def download_directory(self, directory, disk_name): continue subpath = re.sub(directory, '', f) path = os.path.join(disk_name, subpath) - subfolder = re.match('.*\/', path).group(0) + subfolder = re.match('.*\/', path).group(0) # type: ignore if not os.path.exists(subfolder): os.makedirs(subfolder) self.download_to_file(f, path) @clean_object_name - def search(self, pattern, **kwargs): + def search(self, pattern: str, **kwargs) -> None: """Print the objects matching the glob pattern See ``glob`` documentation for details """ matches = self.glob(pattern, verbose=True, **kwargs) - def get_browser(self): + def get_browser(self) -> cottoncandy.browser.BrowserObject: """Return an object which can be tab-completed to browse the contents of the bucket as if it were a file-system @@ -1267,8 +1292,8 @@ def get_browser(self): """ return cottoncandy.browser.S3Directory('', interface = self) - def cp(self, source_name, dest_name, - source_bucket=None, dest_bucket=None, overwrite=False): + def cp(self, source_name: str, dest_name: str, + source_bucket: Optional[str]=None, dest_bucket: Optional[str]=None, overwrite: bool=False): """Copy an object Parameters @@ -1289,8 +1314,8 @@ def cp(self, source_name, dest_name, # TODO: support directories return self.backend_interface.copy(source_name, dest_name, source_bucket, dest_bucket, overwrite) - def mv(self, source_name, dest_name, - source_bucket=None, dest_bucket=None, overwrite=False): + def mv(self, source_name: str, dest_name: str, + source_bucket: Optional[str]=None, dest_bucket: Optional[str]=None, overwrite: bool=False): """Move an object (make copy and delete old object) Parameters @@ -1311,7 +1336,7 @@ def mv(self, source_name, dest_name, # TODO: Support directories return self.backend_interface.move(source_name, dest_name, source_bucket, dest_bucket, overwrite) - def rm(self, object_name, recursive=False, delete=True): + def rm(self, object_name: str, recursive: bool=False, delete: bool=True) -> Any: """Delete an object, or a subtree ('path/to/stuff'). Parameters @@ -1359,7 +1384,7 @@ def rm(self, object_name, recursive=False, delete=True): "nothing found under '%s" print(msg % object_name) - def get_object_owner(self, object_name): + def get_object_owner(self, object_name: str) -> None: self.exists_object(object_name, raise_err=True) ob = self.get_object(object_name) try: @@ -1394,7 +1419,7 @@ def __init__(self, *args, **kwargs): The URL for the S3 gateway force_bucket_creation : bool Create requested bucket if it doesn't exist - backend : 's3'|'gdrive' + backend : 's3'|'gdrive'|'local' which backend to hook on to Returns diff --git a/cottoncandy/localclient.py b/cottoncandy/localclient.py index 3790901..4abdc96 100644 --- a/cottoncandy/localclient.py +++ b/cottoncandy/localclient.py @@ -3,7 +3,7 @@ import os import shutil from io import BytesIO as StringIO -from typing import Optional +from typing import BinaryIO, Optional from .backend import CCBackEnd, CloudStream from .utils import SEPARATOR, remove_root, remove_trivial_magic, sanitize_metadata @@ -15,7 +15,7 @@ class LocalClient(CCBackEnd): """ Client interface for local file system. - Handle metadata in CouldStream objects by storing a json file (.meta.json). + Handle metadata in CloudStream objects by storing a json file (.meta.json). """ def __init__(self, path: str): @@ -23,7 +23,11 @@ def __init__(self, path: str): os.makedirs(path) self.path = path - def check_file_exists(self, cloud_name, bucket_name=None): + @property + def bucket_name(self) -> str: + return self.path + + def check_file_exists(self, cloud_name: str, bucket_name: Optional[str] = None) -> bool: """Checks whether a file exists on the cloud Parameters @@ -42,7 +46,7 @@ def check_file_exists(self, cloud_name, bucket_name=None): bucket_name = self.path return os.path.isfile(os.path.join(bucket_name, cloud_name)) - def upload_stream(self, stream, cloud_name, metadata, permissions, threads = 1): + def upload_stream(self, stream: BinaryIO, cloud_name: str, metadata: dict[str, str], permissions: Optional[str] = None, threads: int = 1) -> None: """Uploads a stream object with a .read() function Parameters @@ -69,7 +73,7 @@ def upload_stream(self, stream, cloud_name, metadata, permissions, threads = 1): with open(metadata_file_name, 'w') as local_file: json.dump(metadata, local_file, indent=4) - def upload_file(self, file_name, cloud_name, permissions, threads = 1): + def upload_file(self, file_name: str, cloud_name: Optional[str] = None, permissions: Optional[str] = None, threads: int = 1) -> None: """Uploads a file from disk Parameters @@ -84,6 +88,8 @@ def upload_file(self, file_name, cloud_name, permissions, threads = 1): ------- bool, upload success """ + if cloud_name is None: + cloud_name = file_name destination = os.path.join(self.path, cloud_name) auto_makedirs(destination) self.copy( @@ -95,7 +101,7 @@ def upload_file(self, file_name, cloud_name, permissions, threads = 1): copy_metadata=False, ) - def download_stream(self, cloud_name, threads = 1): + def download_stream(self, cloud_name: str, threads: int = 1) -> CloudStream: """Downloads a object to an in-memory stream Parameters @@ -122,7 +128,7 @@ def download_stream(self, cloud_name, threads = 1): return CloudStream(content, sanitize_metadata(metadata)) - def download_to_file(self, cloud_name, file_name, threads = 1): + def download_to_file(self, cloud_name: str, file_name: str, threads: int = 1): """Downloads an object directly to disk Parameters @@ -144,7 +150,7 @@ def download_to_file(self, cloud_name, file_name, threads = 1): overwrite=True, ) - def list_directory(self, path, limit): + def list_directory(self, path: str, limit: int) -> list[str]: """Lists the content of a directory Parameters @@ -155,7 +161,7 @@ def list_directory(self, path, limit): Returns ------- - + list[str] """ if (path != '') and (path != '/'): path = remove_root(path) @@ -167,7 +173,7 @@ def list_directory(self, path, limit): results = self._remove_path_and_metadata(results) return results - def list_objects(self, **kwargs): + def list_objects(self, **kwargs) -> list[str]: """Gets all objects contained by backend Returns @@ -189,8 +195,7 @@ def list_objects(self, **kwargs): results = self._remove_path_and_metadata(results) return results - def copy(self, source, destination, source_bucket, destination_bucket, - overwrite, copy_metadata=True): + def copy(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False, copy_metadata: bool = True) -> None: """Copies an object Parameters @@ -225,10 +230,9 @@ def copy(self, source, destination, source_bucket, destination_bucket, shutil.copy(source_metadata, destination_metadata) auto_makedirs(destination) - return shutil.copy(source, destination) + shutil.copy(source, destination) - def move(self, source, destination, source_bucket, destination_bucket, - overwrite): + def move(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False) -> bool: """Moves an object Parameters @@ -261,7 +265,7 @@ def move(self, source, destination, source_bucket, destination_bucket, self._cleanup_empty_dirs(os.path.dirname(source), root_path=source_bucket) return move_result - def delete(self, cloud_name, recursive=False, delete=False): + def delete(self, cloud_name: str, recursive: bool = False, delete: bool = False) -> bool: """Deletes an object Parameters @@ -295,7 +299,7 @@ def delete(self, cloud_name, recursive=False, delete=False): return True @property - def size(self): + def size(self) -> int: """Size of stored cloud items in bytes Returns @@ -312,12 +316,12 @@ def size(self): return total_size - def _remove_path_and_metadata(self, file_list, path=None): + def _remove_path_and_metadata(self, file_list: list[str], path: Optional[str] = None) -> list[str]: """Removes path from filenames, removes .meta.json files from the list. """ if path is None: path = self.path - results = [] + results: list[str] = [] for file_name in file_list: # remove path if file_name.startswith(path): @@ -331,7 +335,7 @@ def _remove_path_and_metadata(self, file_list, path=None): results.append(file_name) return results - def get_object_metadata(self, object_name): + def get_object_metadata(self, object_name: str) -> dict[str, str]: """Get metadata associated with an object""" file_name = os.path.join(self.path, object_name) @@ -349,7 +353,7 @@ def get_object_metadata(self, object_name): return metadata - def get_object_size(self, object_name): + def get_object_size(self, object_name: str) -> int: """Get the size in bytes of an object""" file_name = os.path.join(self.path, object_name) size = os.path.getsize(file_name) diff --git a/cottoncandy/py.typed b/cottoncandy/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/cottoncandy/s3client.py b/cottoncandy/s3client.py index dde27d0..975e4b1 100644 --- a/cottoncandy/s3client.py +++ b/cottoncandy/s3client.py @@ -5,10 +5,12 @@ import os from functools import reduce from io import BytesIO +from typing import Any, BinaryIO, Sequence, Optional, cast from urllib.parse import unquote import boto3 import botocore +import botocore.exceptions # ty wants this explicitly imported from boto3.s3.transfer import TransferConfig from botocore.utils import fix_s3_host from dateutil.tz import tzlocal @@ -43,7 +45,7 @@ class S3Client(CCBackEnd): """ @staticmethod - def connect(ACCESS_KEY, SECRET_KEY, url, **kwargs): + def connect(ACCESS_KEY: str, SECRET_KEY: str, url: str, **kwargs): """Connect to S3 using boto Parameters @@ -66,7 +68,7 @@ def connect(ACCESS_KEY, SECRET_KEY, url, **kwargs): s3.meta.client.meta.events.unregister('before-sign.s3', fix_s3_host) return s3 - def __init__(self, bucket, access_key, secret_key, s3url, force_bucket_creation=False, **kwargs): + def __init__(self, bucket: Optional[str], access_key: str, secret_key: str, s3url: str, force_bucket_creation: bool=False, **kwargs): """Constructor Parameters @@ -81,7 +83,7 @@ def __init__(self, bucket, access_key, secret_key, s3url, force_bucket_creation= self.connection = S3Client.connect(access_key, secret_key, s3url, **kwargs) self.url = s3url - self.bucket_name = None + self._bucket_name: Optional[str] = None if bucket: # bucket given @@ -102,7 +104,7 @@ def __init__(self, bucket, access_key, secret_key, s3url, force_bucket_creation= logging.getLogger('boto3').setLevel(logging.WARNING) logging.getLogger('botocore').setLevel(logging.WARNING) - def get_bucket_name(self, bucket_name): + def get_bucket_name(self, bucket_name: Optional[str] = None) -> Optional[str]: """ Parameters @@ -118,18 +120,23 @@ def get_bucket_name(self, bucket_name): else bucket_name return bucket_name + @property + def bucket_name(self) -> Optional[str]: + return self._bucket_name + @clean_object_name - def check_file_exists(self, object_name, bucket_name=None): + def check_file_exists(self, cloud_name: str, bucket_name: Optional[str] = None) -> bool: """Check whether object exists in bucket Parameters ---------- - object_name : str - The object name - bucket_name + cloud_name : str + The cloud name + bucket_name : str + The bucket name. If None, use the current bucket. """ bucket_name = self.get_bucket_name(bucket_name) - ob = self.connection.Object(key = object_name, bucket_name = bucket_name) + ob = self.connection.Object(key = cloud_name, bucket_name = bucket_name) try: ob.load() @@ -142,7 +149,7 @@ def check_file_exists(self, object_name, bucket_name=None): exists = True return exists - def check_bucket_exists(self, bucket_name): + def check_bucket_exists(self, bucket_name: str): """Check whether the bucket exists Parameters @@ -167,7 +174,7 @@ def check_bucket_exists(self, bucket_name): exists = True return exists - def create_bucket(self, bucket_name, acl=DEFAULT_ACL): + def create_bucket(self, bucket_name: str, acl: str=DEFAULT_ACL): """Create a new bucket Parameters @@ -186,7 +193,7 @@ def create_bucket(self, bucket_name, acl=DEFAULT_ACL): self.connection.create_bucket(Bucket = bucket_name, ACL = acl) self.set_current_bucket(bucket_name) - def set_current_bucket(self, bucket_name): + def set_current_bucket(self, bucket_name: str): """Sets which bucket to use Parameters @@ -199,7 +206,7 @@ def set_current_bucket(self, bucket_name): """ if not self.check_bucket_exists(bucket_name): raise IOError('Bucket "%s" does not exist' % bucket_name) - self.bucket_name = bucket_name + self._bucket_name = bucket_name def get_bucket(self): """Get bucket boto3 object @@ -211,7 +218,7 @@ def get_bucket(self): s3_bucket = self.connection.Bucket(self.bucket_name) return s3_bucket - def list_objects(self, **kwargs): + def list_objects(self, **kwargs) -> Sequence[Any]: """Get list of objects from the bucket This is a wrapper to ``self.get_bucket().bucket.objects`` @@ -239,7 +246,7 @@ def list_objects(self, **kwargs): ) defaults.update(kwargs) bucket = self.get_bucket() - prefix = defaults.pop('filter') + prefix = cast(dict[str, str], defaults.pop('filter')) if prefix['Prefix'] == SEPARATOR: request = bucket.objects @@ -259,10 +266,10 @@ def list_objects(self, **kwargs): return response @property - def size(self): + def size(self) -> int: return self.get_current_bucket_size() - def get_current_bucket_size(self, limit=10 ** 6, page_size=10 ** 6): + def get_current_bucket_size(self, limit: int=10 ** 6, page_size: int=10 ** 6) -> int: """Counts the size of all objects in the current bucket. Parameters @@ -285,6 +292,7 @@ def get_current_bucket_size(self, limit=10 ** 6, page_size=10 ** 6): suspicious round numbers. TODO(anunez): Remove this note when the bug is fixed. """ + assert self.bucket_name is not None, 'Must specify bucket to get size' assert self.check_bucket_exists(self.bucket_name) obs = self.list_objects(limit = limit, page_size = page_size) object_sizes = [t.size for t in obs] @@ -312,7 +320,7 @@ def show_all_buckets(self): print('\n'.join(info)) @clean_object_name - def get_s3_object(self, object_name, bucket_name=None): + def get_s3_object(self, object_name: str, bucket_name: Optional[str] = None): """Get a boto3 object. Create it if it doesn't exist Parameters @@ -327,7 +335,7 @@ def get_s3_object(self, object_name, bucket_name=None): bucket_name = self.get_bucket_name(bucket_name) return self.connection.Object(bucket_name = bucket_name, key = object_name) - def upload_stream(self, stream, cloud_name, metadata, permissions, threads): + def upload_stream(self, stream: BinaryIO, cloud_name: str, metadata: dict[str, str], permissions: Optional[str], threads: int) -> None: """Uploads a stream Parameters @@ -340,6 +348,8 @@ def upload_stream(self, stream, cloud_name, metadata, permissions, threads): ------- """ + assert permissions is not None, 'Permissions must be specified for S3 uploads' + obj = self.get_s3_object(cloud_name) config = TransferConfig(max_concurrency = threads, multipart_chunksize = MPU_CHUNKSIZE, @@ -347,22 +357,22 @@ def upload_stream(self, stream, cloud_name, metadata, permissions, threads): return obj.upload_fileobj(stream, ExtraArgs = {'ACL': permissions, 'Metadata': metadata}, Config = config) - def download_stream(self, object_name, threads): + def download_stream(self, cloud_name: str, threads: int) -> CloudStream: """Download object raw data. This simply calls the object body ``read()`` method. Parameters --------- - object_name : str + cloud_name : str Returns ------- stream file-like stream of object data """ - if not self.check_file_exists(object_name): - raise IOError('Object "%s" does not exist' % object_name) - s3_object = self.get_s3_object(object_name) + if not self.check_file_exists(cloud_name): + raise IOError('Object "%s" does not exist' % cloud_name) + s3_object = self.get_s3_object(cloud_name) config = TransferConfig(max_concurrency = threads, multipart_chunksize = MPU_CHUNKSIZE, multipart_threshold = MPU_THRESHOLD) @@ -371,7 +381,7 @@ def download_stream(self, object_name, threads): byteStream.seek(0) return CloudStream(byteStream, sanitize_metadata(s3_object.metadata)) - def upload_file(self, file_name, cloud_name=None, permissions=DEFAULT_ACL, threads = THREADS): + def upload_file(self, file_name: str, cloud_name: Optional[str] = None, permissions: Optional[str] = None, threads: int = THREADS) -> None: """Upload a file to S3. Parameters @@ -389,6 +399,7 @@ def upload_file(self, file_name, cloud_name=None, permissions=DEFAULT_ACL, threa response : boto3 response """ assert os.path.exists(file_name) + assert permissions is not None, 'Permissions must be specified for S3 uploads' if cloud_name is None: cloud_name = file_name s3_object = self.get_s3_object(cloud_name) @@ -397,24 +408,25 @@ def upload_file(self, file_name, cloud_name=None, permissions=DEFAULT_ACL, threa multipart_threshold = MPU_THRESHOLD) return s3_object.upload_file(file_name, ExtraArgs={'ACL': permissions}, Config = config) - def download_to_file(self, object_name, local_name, threads): + def download_to_file(self, cloud_name: str, local_name: str, threads: int): """Download S3 object to a file Parameters ---------- - object_name : str + cloud_name : str local_name : str Absolute path where the data will be downloaded on disk """ - assert self.check_file_exists(object_name) # make sure object exists - s3_object = self.get_s3_object(object_name) + assert self.check_file_exists(cloud_name) # make sure object exists + s3_object = self.get_s3_object(cloud_name) config = TransferConfig(max_concurrency = threads, multipart_chunksize = MPU_CHUNKSIZE, multipart_threshold = MPU_THRESHOLD) return s3_object.download_file(local_name, Config = config) - def copy(self, source, destination, source_bucket, destination_bucket, overwrite): + def copy(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False) -> Any: source_bucket = self.get_bucket_name(source_bucket) + assert source_bucket is not None, 'Source bucket must be specified' dest_bucket = source_bucket if (destination_bucket is None) else destination_bucket dest_bucket = self.get_bucket_name(dest_bucket) @@ -428,13 +440,13 @@ def copy(self, source, destination, source_bucket, destination_bucket, overwrite ob_new.copy_from(CopySource = fpath) return ob_new - def move(self, source, destination, source_bucket, destination_bucket, overwrite): + def move(self, source: str, destination: str, source_bucket: Optional[str] = None, destination_bucket: Optional[str] = None, overwrite: bool = False) -> Any: new_ob = self.copy(source, destination, source_bucket, destination_bucket, overwrite) old_ob = self.get_s3_object(source, bucket_name = source_bucket) old_ob.delete() return new_ob - def list_directory(self, path, limit): + def list_directory(self, path: str, limit: int) -> list[str]: """List the contents of a "directory" Parameters @@ -458,26 +470,26 @@ def list_directory(self, path, limit): Delimiter = SEPARATOR, Prefix = path, MaxKeys = limit) - object_names = [] + object_names: list[str] = [] if 'CommonPrefixes' in response: # we got common paths - object_list = [list(t.values()) for t in response['CommonPrefixes']] - object_names += reduce(lambda x, y: x + y, object_list) + object_list: list[list[str]] = [list(t.values()) for t in response['CommonPrefixes']] + object_names = object_names + list(reduce(lambda x, y: x + y, object_list)) # type: ignore[operator] if 'Contents' in response: # we got objects on the leaf nodes object_names += unquote_names([t['Key'] for t in response['Contents']]) return [os.path.normpath(n) for n in object_names] - def delete(self, object_name, recursive=False, delete=False): + def delete(self, cloud_name: str, recursive: bool=False, delete: bool=False): raise RuntimeError('Deleting on S3 backend is implemented by cottoncandy interface object') - def get_object_metadata(self, object_name): + def get_object_metadata(self, object_name: str) -> dict[str, str]: """Get metadata associated with an object""" s3_object = self.get_s3_object(object_name) metadata = sanitize_metadata(s3_object.metadata) return metadata - def get_object_size(self, object_name): + def get_object_size(self, object_name: str) -> int: """Get the size in bytes of an object""" s3_object = self.get_s3_object(object_name) size = s3_object.content_length diff --git a/cottoncandy/tests/conftest.py b/cottoncandy/tests/conftest.py index cae6f0c..8a60a59 100644 --- a/cottoncandy/tests/conftest.py +++ b/cottoncandy/tests/conftest.py @@ -57,5 +57,8 @@ def cci(request): cci.wait_time = 0.001 yield cci + else: + raise ValueError("Invalid client type") + # cleanup the directory entirely cci.rm(directory, recursive=True) diff --git a/cottoncandy/tests/test_compression.py b/cottoncandy/tests/test_compression.py index 6690154..7b57f69 100644 --- a/cottoncandy/tests/test_compression.py +++ b/cottoncandy/tests/test_compression.py @@ -1,4 +1,5 @@ import time +from typing import Literal import numpy as np @@ -6,7 +7,7 @@ def content_generator(): size_mb = 101 - orders = ['F', 'C'] + orders: list[Literal['C', 'F']] = ['F', 'C'] types = [ 'float64', ] diff --git a/cottoncandy/tests/test_roundtrip.py b/cottoncandy/tests/test_roundtrip.py index 5b12eca..d3b5832 100644 --- a/cottoncandy/tests/test_roundtrip.py +++ b/cottoncandy/tests/test_roundtrip.py @@ -1,12 +1,13 @@ import os import tempfile import time +from typing import Literal import numpy as np def content_generator(): - orders = ['C', 'F'] + orders: list[Literal['C', 'F']] = ['C', 'F'] types = [ 'float16', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64', 'uint8', 'uint16', 'uint32', 'int', 'float' diff --git a/cottoncandy/tests/test_roundtrip_big.py b/cottoncandy/tests/test_roundtrip_big.py index 60dfbfc..87a427a 100644 --- a/cottoncandy/tests/test_roundtrip_big.py +++ b/cottoncandy/tests/test_roundtrip_big.py @@ -1,11 +1,12 @@ import time +from typing import Literal import numpy as np def content_generator(): size_mb = 200 - orders = ['C', 'F'] + orders: list[Literal['C', 'F']] = ['C', 'F'] types = [ 'float64', ] diff --git a/cottoncandy/utils.py b/cottoncandy/utils.py index bfadb6e..dffbe38 100644 --- a/cottoncandy/utils.py +++ b/cottoncandy/utils.py @@ -1,14 +1,18 @@ '''Helper functions ''' +from io import BytesIO import itertools import os import re -import string import zlib from functools import wraps +from typing import Iterable, Iterator, Optional, cast, Any, BinaryIO, Callable, TypeVar, Union + + from urllib.parse import unquote import numpy as np +import numpy.typing as npt import six from dateutil.tz import tzlocal @@ -44,14 +48,14 @@ # misc functions ############################## -def sanitize_metadata(metadict): +def sanitize_metadata(metadict: dict[str, str]) -> dict[str, str]: outdict = {} for key,val in metadict.items(): outdict[key.lower()] = val return outdict -def pathjoin(a, *p): +def pathjoin(a: str, *p: str) -> str: """Join two or more pathname components, inserting SEPARATOR as needed. If any component is an absolute path, all previous path components will be discarded. An empty last part will result in a path that @@ -67,10 +71,10 @@ def pathjoin(a, *p): return path -def string2bool(mstring): +def string2bool(mstring: str) -> Union[bool, None]: ''' ''' - truth_value = False + truth_value: Union[bool, None] = False if mstring in ['True','true', 'tru', 't', 'y','yes', '1']: truth_value = True @@ -79,7 +83,7 @@ def string2bool(mstring): return truth_value -def bytes2human(nbytes): +def bytes2human(nbytes: int) -> str: '''Return string representation of bytes. Parameters @@ -127,7 +131,7 @@ def get_object_size(boto_s3_object): return boto_s3_object.meta.data['ContentLength']/2.**20 -def get_fileobject_size(file_object): +def get_fileobject_size(file_object: BinaryIO) -> int: '''Return byte size of file-object Parameters @@ -217,7 +221,7 @@ def objects2names(objects): return [unquote(t.key) for t in objects] -def unquote_names(object_names): +def unquote_names(object_names: list[str]) -> list[str]: '''Clean URL names from a list. Parameters @@ -253,15 +257,19 @@ def print_objects(object_list): ############################## -def clean_object_name(input_function): +# Decorator typing from mypy docs: +# https://mypy.readthedocs.io/en/stable/generics.html#declaring-decorators +F = TypeVar('F', bound=Callable[..., Any]) + +def clean_object_name(input_function: F) -> F: '''Remove leading "/" from object_name This is important for compatibility with S3fs. S3fs does not list objects with a "/" prefix. ''' @wraps(input_function) - def iremove_root(self, object_name, *args, **kwargs): - object_name = re.sub('//+', '/', object_name) + def iremove_root(self, cloud_name: str, *args: Any, **kwargs: Any) -> Any: + object_name = re.sub('//+', '/', cloud_name) if object_name == '': pass @@ -269,10 +277,10 @@ def iremove_root(self, object_name, *args, **kwargs): object_name = object_name[1:] return input_function(self, object_name, *args, **kwargs) - return iremove_root + return cast(F, iremove_root) -def remove_root(string_): +def remove_root(string_: str) -> str: '''remove leading "/" from a string''' if string_[0] == SEPARATOR: string_ = string_[1:] @@ -292,7 +300,7 @@ def has_start_digit(s): -def has_magic(s): +def has_magic(s: str) -> bool: '''Check string to see if it has any glob magic ''' return MAGIC_CHECK.search(s) is not None @@ -312,13 +320,13 @@ def has_trivial_magic(s): return False -def has_real_magic(s): +def has_real_magic(s: str) -> bool: '''Check if string has non-trivial glob pattern ''' return has_magic(s) and (not has_trivial_magic(s)) -def remove_trivial_magic(s): +def remove_trivial_magic(s: str) -> str: ''' * xxx/* -> xxx/ * xxx/ -> xxx/ @@ -330,7 +338,7 @@ def remove_trivial_magic(s): return s[:-1] # remove '*' at end -def split_uri(uri, pattern='s3://', separator='/'): +def split_uri(uri: str, pattern: str='s3://', separator: str='/') -> tuple[str, str]: """Convert a URI to a bucket, object name tuple. 's3://bucket/path/to/thing' -> ('bucket', 'path/to/thing') @@ -342,7 +350,7 @@ def split_uri(uri, pattern='s3://', separator='/'): return bucket, path -def mk_aws_path(path): +def mk_aws_path(path: str) -> str: """Make the `path` behave as expected when querying S3 with `list_objects`. @@ -367,7 +375,7 @@ def mk_aws_path(path): ############################## -def generate_ndarray_chunks(arr, axis=None, buffersize=100*MB): +def generate_ndarray_chunks(arr: npt.NDArray, axis: Optional[int]=None, buffersize: int=100*MB) -> Iterator[tuple[tuple[int, ...], npt.NDArray]]: '''A generator that splits an array into chunks of desired byte size Parameters @@ -414,7 +422,7 @@ def generate_ndarray_chunks(arr, axis=None, buffersize=100*MB): logii = ((np.log(buffersize) - np.log(arr.itemsize)) - logsum)/factor ii = int(np.ceil(np.exp(logii))) - dim_nchunks = map(lambda x: int(np.ceil(x/ii)) + 1, shape) + dim_nchunks: Iterable[int] = map(lambda x: int(np.ceil(x/ii)) + 1, shape) if axis is not None: # only slicing one dimension @@ -433,7 +441,7 @@ def generate_ndarray_chunks(arr, axis=None, buffersize=100*MB): yield chunk_coords, arr[slicers] -def read_buffered(frm, to, buffersize=64): +def read_buffered(frm: BinaryIO, to: npt.NDArray, buffersize: int = 64) -> None: '''Fill a numpy n-d array with file-like object contents Parameters @@ -450,7 +458,8 @@ def read_buffered(frm, to, buffersize=64): else: vw = to.view() vw.shape = (-1,) # Must be a ravel-able object - vw.dtype = np.dtype('uint8') # 256 values in each byte + # 256 values in each byte + vw.dtype = np.dtype('uint8') # type: ignore[misc] for ci in range(int(np.ceil(nbytes_total / float(buffersize)))): start = ci * buffersize @@ -460,10 +469,10 @@ def read_buffered(frm, to, buffersize=64): elif six.PY3: vw.data[start:end] = frm.read(end - start) else: - raise("Unknown python version") # not sure six will ever do anything here (6=2x3) + raise Exception("Unknown python version") # not sure six will ever do anything here (6=2x3) -class GzipInputStream: +class GzipInputStream(BytesIO): """Simple class that allow streaming reads from GZip files (from https://gist.github.com/beaufour/4205533). @@ -473,22 +482,23 @@ class GzipInputStream: Adapted from: http://effbot.org/librarybook/zlib-example-4.py """ - def __init__(self, fileobj, block_size=16384): + def __init__(self, fileobj: BinaryIO, block_size: int=16384): """ Initialize with the given file-like object. @param fileobj: file-like object, """ + super().__init__() self.BLOCK_SIZE = block_size # Read block size # zlib window buffer size, set to gzip's format self.WINDOW_BUFFER_SIZE = 16 + zlib.MAX_WBITS self._file = fileobj - self._zip = zlib.decompressobj(self.WINDOW_BUFFER_SIZE) + self._zip: Optional[zlib._Decompress] = zlib.decompressobj(self.WINDOW_BUFFER_SIZE) self._offset = 0 # position in unzipped stream self._data = b'' - def __fill(self, num_bytes): + def __fill(self, num_bytes: int) -> None: """ Fill the internal buffer with 'num_bytes' of data. @@ -510,7 +520,7 @@ def __fill(self, num_bytes): def __iter__(self): return self - def seek(self, offset, whence=0): + def seek(self, offset: int, whence: int = 0) -> int: if whence == 0: position = offset elif whence == 1: @@ -525,10 +535,13 @@ def seek(self, offset, whence=0): if not self.read(min(position - self._offset, self.BLOCK_SIZE)): break - def tell(self): + return position + + def tell(self) -> int: return self._offset - def read(self, size=0): + def read(self, size: Optional[int] = 0) -> bytes: + assert size is not None self.__fill(size) if size: data = self._data[:size] @@ -539,23 +552,25 @@ def read(self, size=0): self._offset = self._offset + len(data) return data - def next(self): + def next(self) -> bytes: line = self.readline() if not line: raise StopIteration() return line - def readline(self): + def readline(self, size: Optional[int] = None) -> bytes: + assert size is None # make sure we have an entire line - while self._zip and "\n" not in self._data: + while self._zip and b"\n" not in self._data: self.__fill(len(self._data) + 512) - pos = string.find(self._data, "\n") + 1 + pos = self._data.find(b"\n") + 1 if pos <= 0: return self.read() return self.read(pos) - def readlines(self): + def readlines(self, size: Optional[int] = None) -> list[bytes]: + assert size is None lines = [] while True: line = self.readline() diff --git a/pyproject.toml b/pyproject.toml index 1d14eae..d139ebd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "numcodecs>=0.5.5", "urllib3>=1.26", # newer urllib3 is required but unspecified by botocore ] -requires-python = ">=3.7" +requires-python = ">=3.9" [project.optional-dependencies] extra = [ @@ -38,6 +38,8 @@ extra = [ test = [ "codecov", "flake8", + "ipython", # IPython is needed for something in gdriveclient + "mypy", "pytest", "pytest-cov", "pytest-rerunfailures", @@ -62,7 +64,7 @@ Documentation = "http://gallantlab.github.io/cottoncandy/" "Bug Tracker" = "https://github.com/gallantlab/cottoncandy/issues" [tool.pytest.ini_options] -addopts = "--cov=. --cov-report xml:coverage.xml" +addopts = "--cov=. --cov-report xml" [tool.ruff] exclude = ["appdirs.py", "*.ipynb"] @@ -73,7 +75,7 @@ packages = ["cottoncandy"] include-package-data = true [tool.setuptools.package-data] -cottoncandy = ["defaults.cfg"] +cottoncandy = ["defaults.cfg", "py.typed"] [tool.setuptools.dynamic] version = {attr = "cottoncandy.__version__"} @@ -84,3 +86,7 @@ skip = '.git*,*.svg,*.css,*.min.*,docs,*.map,.npm,.cache' check-hidden = true ignore-regex = '^\s*"image/\S+": ".*' # ignore-words-list = '' + +[tool.mypy] +allow_redefinition = true +disable_error_code = "import-untyped"