Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/fsspec/utils.py: 14%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1from __future__ import annotations
3import contextlib
4import logging
5import math
6import os
7import re
8import sys
9import tempfile
10from collections.abc import Iterable, Iterator, Sequence
11from functools import partial
12from hashlib import md5
13from importlib.metadata import version
14from typing import (
15 IO,
16 TYPE_CHECKING,
17 Any,
18 Callable,
19 TypeVar,
20)
21from urllib.parse import urlsplit
23if TYPE_CHECKING:
24 import pathlib
26 from typing_extensions import TypeGuard
28 from fsspec.spec import AbstractFileSystem
31DEFAULT_BLOCK_SIZE = 5 * 2**20
33T = TypeVar("T")
36def infer_storage_options(
37 urlpath: str, inherit_storage_options: dict[str, Any] | None = None
38) -> dict[str, Any]:
39 """Infer storage options from URL path and merge it with existing storage
40 options.
42 Parameters
43 ----------
44 urlpath: str or unicode
45 Either local absolute file path or URL (hdfs://namenode:8020/file.csv)
46 inherit_storage_options: dict (optional)
47 Its contents will get merged with the inferred information from the
48 given path
50 Returns
51 -------
52 Storage options dict.
54 Examples
55 --------
56 >>> infer_storage_options('/mnt/datasets/test.csv') # doctest: +SKIP
57 {"protocol": "file", "path", "/mnt/datasets/test.csv"}
58 >>> infer_storage_options(
59 ... 'hdfs://username:pwd@node:123/mnt/datasets/test.csv?q=1',
60 ... inherit_storage_options={'extra': 'value'},
61 ... ) # doctest: +SKIP
62 {"protocol": "hdfs", "username": "username", "password": "pwd",
63 "host": "node", "port": 123, "path": "/mnt/datasets/test.csv",
64 "url_query": "q=1", "extra": "value"}
65 """
66 # Handle Windows paths including disk name in this special case
67 if (
68 re.match(r"^[a-zA-Z]:[\\/]", urlpath)
69 or re.match(r"^[a-zA-Z0-9]+://", urlpath) is None
70 ):
71 return {"protocol": "file", "path": urlpath}
73 parsed_path = urlsplit(urlpath)
74 protocol = parsed_path.scheme or "file"
75 if parsed_path.fragment:
76 path = "#".join([parsed_path.path, parsed_path.fragment])
77 else:
78 path = parsed_path.path
79 if protocol == "file":
80 # Special case parsing file protocol URL on Windows according to:
81 # https://msdn.microsoft.com/en-us/library/jj710207.aspx
82 windows_path = re.match(r"^/([a-zA-Z])[:|]([\\/].*)$", path)
83 if windows_path:
84 drive, path = windows_path.groups()
85 path = f"{drive}:{path}"
87 if protocol in ["http", "https"]:
88 # for HTTP, we don't want to parse, as requests will anyway
89 return {"protocol": protocol, "path": urlpath}
91 options: dict[str, Any] = {"protocol": protocol, "path": path}
93 if parsed_path.netloc:
94 # Parse `hostname` from netloc manually because `parsed_path.hostname`
95 # lowercases the hostname which is not always desirable (e.g. in S3):
96 # https://github.com/dask/dask/issues/1417
97 options["host"] = parsed_path.netloc.rsplit("@", 1)[-1].rsplit(":", 1)[0]
99 if protocol in ("s3", "s3a", "gcs", "gs"):
100 options["path"] = options["host"] + options["path"]
101 else:
102 options["host"] = options["host"]
103 if parsed_path.port:
104 options["port"] = parsed_path.port
105 if parsed_path.username:
106 options["username"] = parsed_path.username
107 if parsed_path.password:
108 options["password"] = parsed_path.password
110 if parsed_path.query:
111 options["url_query"] = parsed_path.query
112 if parsed_path.fragment:
113 options["url_fragment"] = parsed_path.fragment
115 if inherit_storage_options:
116 update_storage_options(options, inherit_storage_options)
118 return options
121def update_storage_options(
122 options: dict[str, Any], inherited: dict[str, Any] | None = None
123) -> None:
124 if not inherited:
125 inherited = {}
126 collisions = set(options) & set(inherited)
127 if collisions:
128 for collision in collisions:
129 if options.get(collision) != inherited.get(collision):
130 raise KeyError(
131 f"Collision between inferred and specified storage "
132 f"option:\n{collision}"
133 )
134 options.update(inherited)
137# Compression extensions registered via fsspec.compression.register_compression
138compressions: dict[str, str] = {}
141def infer_compression(filename: str) -> str | None:
142 """Infer compression, if available, from filename.
144 Infer a named compression type, if registered and available, from filename
145 extension. This includes builtin (gz, bz2, zip) compressions, as well as
146 optional compressions. See fsspec.compression.register_compression.
147 """
148 extension = os.path.splitext(filename)[-1].strip(".").lower()
149 if extension in compressions:
150 return compressions[extension]
151 return None
154def build_name_function(max_int: float) -> Callable[[int], str]:
155 """Returns a function that receives a single integer
156 and returns it as a string padded by enough zero characters
157 to align with maximum possible integer
159 >>> name_f = build_name_function(57)
161 >>> name_f(7)
162 '07'
163 >>> name_f(31)
164 '31'
165 >>> build_name_function(1000)(42)
166 '0042'
167 >>> build_name_function(999)(42)
168 '042'
169 >>> build_name_function(0)(0)
170 '0'
171 """
172 # handle corner cases max_int is 0 or exact power of 10
173 max_int += 1e-8
175 pad_length = int(math.ceil(math.log10(max_int)))
177 def name_function(i: int) -> str:
178 return str(i).zfill(pad_length)
180 return name_function
183def seek_delimiter(file: IO[bytes], delimiter: bytes, blocksize: int) -> bool:
184 r"""Seek current file to file start, file end, or byte after delimiter seq.
186 Seeks file to next chunk delimiter, where chunks are defined on file start,
187 a delimiting sequence, and file end. Use file.tell() to see location afterwards.
188 Note that file start is a valid split, so must be at offset > 0 to seek for
189 delimiter.
191 Parameters
192 ----------
193 file: a file
194 delimiter: bytes
195 a delimiter like ``b'\n'`` or message sentinel, matching file .read() type
196 blocksize: int
197 Number of bytes to read from the file at once.
200 Returns
201 -------
202 Returns True if a delimiter was found, False if at file start or end.
204 """
206 if file.tell() == 0:
207 # beginning-of-file, return without seek
208 return False
210 # Interface is for binary IO, with delimiter as bytes, but initialize last
211 # with result of file.read to preserve compatibility with text IO.
212 last: bytes | None = None
213 while True:
214 current = file.read(blocksize)
215 if not current:
216 # end-of-file without delimiter
217 return False
218 full = last + current if last else current
219 try:
220 if delimiter in full:
221 i = full.index(delimiter)
222 file.seek(file.tell() - (len(full) - i) + len(delimiter))
223 return True
224 elif len(current) < blocksize:
225 # end-of-file without delimiter
226 return False
227 except (OSError, ValueError):
228 pass
229 last = full[-len(delimiter) :]
232def read_block(
233 f: IO[bytes],
234 offset: int,
235 length: int | None,
236 delimiter: bytes | None = None,
237 split_before: bool = False,
238) -> bytes:
239 """Read a block of bytes from a file
241 Parameters
242 ----------
243 f: File
244 Open file
245 offset: int
246 Byte offset to start read
247 length: int
248 Number of bytes to read, read through end of file if None
249 delimiter: bytes (optional)
250 Ensure reading starts and stops at delimiter bytestring
251 split_before: bool (optional)
252 Start/stop read *before* delimiter bytestring.
255 If using the ``delimiter=`` keyword argument we ensure that the read
256 starts and stops at delimiter boundaries that follow the locations
257 ``offset`` and ``offset + length``. If ``offset`` is zero then we
258 start at zero, regardless of delimiter. The bytestring returned WILL
259 include the terminating delimiter string.
261 Examples
262 --------
264 >>> from io import BytesIO # doctest: +SKIP
265 >>> f = BytesIO(b'Alice, 100\\nBob, 200\\nCharlie, 300') # doctest: +SKIP
266 >>> read_block(f, 0, 13) # doctest: +SKIP
267 b'Alice, 100\\nBo'
269 >>> read_block(f, 0, 13, delimiter=b'\\n') # doctest: +SKIP
270 b'Alice, 100\\nBob, 200\\n'
272 >>> read_block(f, 10, 10, delimiter=b'\\n') # doctest: +SKIP
273 b'Bob, 200\\nCharlie, 300'
274 """
275 if delimiter:
276 f.seek(offset)
277 found_start_delim = seek_delimiter(f, delimiter, 2**16)
278 if length is None:
279 return f.read()
280 start = f.tell()
281 length -= start - offset
283 f.seek(start + length)
284 found_end_delim = seek_delimiter(f, delimiter, 2**16)
285 end = f.tell()
287 # Adjust split location to before delimiter if seek found the
288 # delimiter sequence, not start or end of file.
289 if found_start_delim and split_before:
290 start -= len(delimiter)
292 if found_end_delim and split_before:
293 end -= len(delimiter)
295 offset = start
296 length = end - start
298 f.seek(offset)
300 # TODO: allow length to be None and read to the end of the file?
301 assert length is not None
302 b = f.read(length)
303 return b
306def tokenize(*args: Any, **kwargs: Any) -> str:
307 """Deterministic token
309 (modified from dask.base)
311 >>> tokenize([1, 2, '3'])
312 '9d71491b50023b06fc76928e6eddb952'
314 >>> tokenize('Hello') == tokenize('Hello')
315 True
316 """
317 if kwargs:
318 args += (kwargs,)
319 try:
320 h = md5(str(args).encode())
321 except ValueError:
322 # FIPS systems: https://github.com/fsspec/filesystem_spec/issues/380
323 h = md5(str(args).encode(), usedforsecurity=False)
324 return h.hexdigest()
327def stringify_path(filepath: str | os.PathLike[str] | pathlib.Path) -> str:
328 """Attempt to convert a path-like object to a string.
330 Parameters
331 ----------
332 filepath: object to be converted
334 Returns
335 -------
336 filepath_str: maybe a string version of the object
338 Notes
339 -----
340 Objects supporting the fspath protocol are coerced according to its
341 __fspath__ method.
343 For backwards compatibility with older Python version, pathlib.Path
344 objects are specially coerced.
346 Any other object is passed through unchanged, which includes bytes,
347 strings, buffers, or anything else that's not even path-like.
348 """
349 if isinstance(filepath, str):
350 return filepath
351 elif hasattr(filepath, "__fspath__"):
352 return filepath.__fspath__()
353 elif hasattr(filepath, "path"):
354 return filepath.path
355 else:
356 return filepath # type: ignore[return-value]
359def make_instance(
360 cls: Callable[..., T], args: Sequence[Any], kwargs: dict[str, Any]
361) -> T:
362 inst = cls(*args, **kwargs)
363 inst._determine_worker() # type: ignore[attr-defined]
364 return inst
367def common_prefix(paths: Iterable[str]) -> str:
368 """For a list of paths, find the shortest prefix common to all"""
369 parts = [p.split("/") for p in paths]
370 lmax = min(len(p) for p in parts)
371 end = 0
372 for i in range(lmax):
373 end = all(p[i] == parts[0][i] for p in parts)
374 if not end:
375 break
376 i += end
377 return "/".join(parts[0][:i])
380def other_paths(
381 paths: list[str],
382 path2: str | list[str],
383 exists: bool = False,
384 flatten: bool = False,
385) -> list[str]:
386 """In bulk file operations, construct a new file tree from a list of files
388 Parameters
389 ----------
390 paths: list of str
391 The input file tree
392 path2: str or list of str
393 Root to construct the new list in. If this is already a list of str, we just
394 assert it has the right number of elements.
395 exists: bool (optional)
396 For a str destination, it is already exists (and is a dir), files should
397 end up inside.
398 flatten: bool (optional)
399 Whether to flatten the input directory tree structure so that the output files
400 are in the same directory.
402 Returns
403 -------
404 list of str
405 """
407 if isinstance(path2, str):
408 path2 = path2.rstrip("/")
410 if flatten:
411 path2 = ["/".join((path2, p.split("/")[-1])) for p in paths]
412 else:
413 cp = common_prefix(paths)
414 if exists:
415 cp = cp.rsplit("/", 1)[0]
416 if not cp and all(not s.startswith("/") for s in paths):
417 path2 = ["/".join([path2, p]) for p in paths]
418 else:
419 path2 = [p.replace(cp, path2, 1) for p in paths]
420 else:
421 assert len(paths) == len(path2)
422 return path2
425def is_exception(obj: Any) -> bool:
426 return isinstance(obj, BaseException)
429def isfilelike(f: Any) -> TypeGuard[IO[bytes]]:
430 return all(hasattr(f, attr) for attr in ["read", "close", "tell"])
433def get_protocol(url: str) -> str:
434 url = stringify_path(url)
435 parts = re.split(r"(\:\:|\://)", url, maxsplit=1)
436 if len(parts) > 1:
437 return parts[0]
438 return "file"
441def get_file_extension(url: str) -> str:
442 url = stringify_path(url)
443 ext_parts = url.rsplit(".", 1)
444 if len(ext_parts) > 1:
445 return ext_parts[-1]
446 return ""
449def can_be_local(path: str) -> bool:
450 """Can the given URL be used with open_local?"""
451 from fsspec import get_filesystem_class
453 try:
454 return getattr(get_filesystem_class(get_protocol(path)), "local_file", False)
455 except (ValueError, ImportError):
456 # not in registry or import failed
457 return False
460def get_package_version_without_import(name: str) -> str | None:
461 """For given package name, try to find the version without importing it
463 Import and package.__version__ is still the backup here, so an import
464 *might* happen.
466 Returns either the version string, or None if the package
467 or the version was not readily found.
468 """
469 if name in sys.modules:
470 mod = sys.modules[name]
471 if hasattr(mod, "__version__"):
472 return mod.__version__
473 try:
474 return version(name)
475 except: # noqa: E722
476 pass
477 try:
478 import importlib
480 mod = importlib.import_module(name)
481 return mod.__version__
482 except (ImportError, AttributeError):
483 return None
486def setup_logging(
487 logger: logging.Logger | None = None,
488 logger_name: str | None = None,
489 level: str = "DEBUG",
490 clear: bool = True,
491) -> logging.Logger:
492 if logger is None and logger_name is None:
493 raise ValueError("Provide either logger object or logger name")
494 logger = logger or logging.getLogger(logger_name)
495 handle = logging.StreamHandler()
496 formatter = logging.Formatter(
497 "%(asctime)s - %(name)s - %(levelname)s - %(funcName)s -- %(message)s"
498 )
499 handle.setFormatter(formatter)
500 if clear:
501 logger.handlers.clear()
502 logger.addHandler(handle)
503 logger.setLevel(level)
504 return logger
507def _unstrip_protocol(name: str, fs: AbstractFileSystem) -> str:
508 return fs.unstrip_protocol(name)
511def mirror_from(
512 origin_name: str, methods: Iterable[str]
513) -> Callable[[type[T]], type[T]]:
514 """Mirror attributes and methods from the given
515 origin_name attribute of the instance to the
516 decorated class"""
518 def origin_getter(method: str, self: Any) -> Any:
519 origin = getattr(self, origin_name)
520 return getattr(origin, method)
522 def wrapper(cls: type[T]) -> type[T]:
523 for method in methods:
524 wrapped_method = partial(origin_getter, method)
525 setattr(cls, method, property(wrapped_method))
526 return cls
528 return wrapper
531@contextlib.contextmanager
532def nullcontext(obj: T) -> Iterator[T]:
533 yield obj
536def merge_offset_ranges(
537 paths: list[str],
538 starts: list[int] | int,
539 ends: list[int] | int,
540 max_gap: int = 0,
541 max_block: int | None = None,
542 sort: bool = True,
543) -> tuple[list[str], list[int], list[int]]:
544 """Merge adjacent byte-offset ranges when the inter-range
545 gap is <= `max_gap`, and when the merged byte range does not
546 exceed `max_block` (if specified). By default, this function
547 will re-order the input paths and byte ranges to ensure sorted
548 order. If the user can guarantee that the inputs are already
549 sorted, passing `sort=False` will skip the re-ordering.
550 """
551 # Check input
552 if not isinstance(paths, list):
553 raise TypeError
554 if not isinstance(starts, list):
555 starts = [starts] * len(paths)
556 if not isinstance(ends, list):
557 ends = [ends] * len(paths)
558 if len(starts) != len(paths) or len(ends) != len(paths):
559 raise ValueError
561 # Early Return
562 if len(starts) <= 1:
563 return paths, starts, ends
565 starts = [s or 0 for s in starts]
566 # Sort by paths and then ranges if `sort=True`
567 if sort:
568 paths, starts, ends = (
569 list(v)
570 for v in zip(
571 *sorted(
572 zip(paths, starts, ends),
573 )
574 )
575 )
577 if paths:
578 # Loop through the coupled `paths`, `starts`, and
579 # `ends`, and merge adjacent blocks when appropriate
580 new_paths = paths[:1]
581 new_starts = starts[:1]
582 new_ends = ends[:1]
583 for i in range(1, len(paths)):
584 if paths[i] == paths[i - 1] and new_ends[-1] is None:
585 continue
586 elif (
587 paths[i] != paths[i - 1]
588 or ((starts[i] - new_ends[-1]) > max_gap)
589 or (max_block is not None and (ends[i] - new_starts[-1]) > max_block)
590 ):
591 # Cannot merge with previous block.
592 # Add new `paths`, `starts`, and `ends` elements
593 new_paths.append(paths[i])
594 new_starts.append(starts[i])
595 new_ends.append(ends[i])
596 else:
597 # Merge with previous block by updating the
598 # last element of `ends`
599 new_ends[-1] = ends[i]
600 return new_paths, new_starts, new_ends
602 # `paths` is empty. Just return input lists
603 return paths, starts, ends
606def file_size(filelike: IO[bytes]) -> int:
607 """Find length of any open read-mode file-like"""
608 pos = filelike.tell()
609 try:
610 return filelike.seek(0, 2)
611 finally:
612 filelike.seek(pos)
615@contextlib.contextmanager
616def atomic_write(path: str, mode: str = "wb"):
617 """
618 A context manager that opens a temporary file next to `path` and, on exit,
619 replaces `path` with the temporary file, thereby updating `path`
620 atomically.
621 """
622 fd, fn = tempfile.mkstemp(
623 dir=os.path.dirname(path), prefix=os.path.basename(path) + "-"
624 )
625 try:
626 with open(fd, mode) as fp:
627 yield fp
628 except BaseException:
629 with contextlib.suppress(FileNotFoundError):
630 os.unlink(fn)
631 raise
632 else:
633 os.replace(fn, path)
636def _translate(pat, STAR, QUESTION_MARK):
637 # Copied from: https://github.com/python/cpython/pull/106703.
638 res: list[str] = []
639 add = res.append
640 i, n = 0, len(pat)
641 while i < n:
642 c = pat[i]
643 i = i + 1
644 if c == "*":
645 # compress consecutive `*` into one
646 if (not res) or res[-1] is not STAR:
647 add(STAR)
648 elif c == "?":
649 add(QUESTION_MARK)
650 elif c == "[":
651 j = i
652 if j < n and pat[j] == "!":
653 j = j + 1
654 if j < n and pat[j] == "]":
655 j = j + 1
656 while j < n and pat[j] != "]":
657 j = j + 1
658 if j >= n:
659 add("\\[")
660 else:
661 stuff = pat[i:j]
662 if "-" not in stuff:
663 stuff = stuff.replace("\\", r"\\")
664 else:
665 chunks = []
666 k = i + 2 if pat[i] == "!" else i + 1
667 while True:
668 k = pat.find("-", k, j)
669 if k < 0:
670 break
671 chunks.append(pat[i:k])
672 i = k + 1
673 k = k + 3
674 chunk = pat[i:j]
675 if chunk:
676 chunks.append(chunk)
677 else:
678 chunks[-1] += "-"
679 # Remove empty ranges -- invalid in RE.
680 for k in range(len(chunks) - 1, 0, -1):
681 if chunks[k - 1][-1] > chunks[k][0]:
682 chunks[k - 1] = chunks[k - 1][:-1] + chunks[k][1:]
683 del chunks[k]
684 # Escape backslashes and hyphens for set difference (--).
685 # Hyphens that create ranges shouldn't be escaped.
686 stuff = "-".join(
687 s.replace("\\", r"\\").replace("-", r"\-") for s in chunks
688 )
689 # Escape set operations (&&, ~~ and ||).
690 stuff = re.sub(r"([&~|])", r"\\\1", stuff)
691 i = j + 1
692 if not stuff:
693 # Empty range: never match.
694 add("(?!)")
695 elif stuff == "!":
696 # Negated empty range: match any character.
697 add(".")
698 else:
699 if stuff[0] == "!":
700 stuff = "^" + stuff[1:]
701 elif stuff[0] in ("^", "["):
702 stuff = "\\" + stuff
703 add(f"[{stuff}]")
704 else:
705 add(re.escape(c))
706 assert i == n
707 return res
710def glob_translate(pat):
711 # Copied from: https://github.com/python/cpython/pull/106703.
712 # The keyword parameters' values are fixed to:
713 # recursive=True, include_hidden=True, seps=None
714 """Translate a pathname with shell wildcards to a regular expression."""
715 if os.path.altsep:
716 seps = os.path.sep + os.path.altsep
717 else:
718 seps = os.path.sep
719 escaped_seps = "".join(map(re.escape, seps))
720 any_sep = f"[{escaped_seps}]" if len(seps) > 1 else escaped_seps
721 not_sep = f"[^{escaped_seps}]"
722 one_last_segment = f"{not_sep}+"
723 one_segment = f"{one_last_segment}{any_sep}"
724 any_segments = f"(?:.+{any_sep})?"
725 any_last_segments = ".*"
726 results = []
727 parts = re.split(any_sep, pat)
728 last_part_idx = len(parts) - 1
729 for idx, part in enumerate(parts):
730 if part == "*":
731 results.append(one_segment if idx < last_part_idx else one_last_segment)
732 continue
733 if part == "**":
734 results.append(any_segments if idx < last_part_idx else any_last_segments)
735 continue
736 elif "**" in part:
737 raise ValueError(
738 "Invalid pattern: '**' can only be an entire path component"
739 )
740 if part:
741 results.extend(_translate(part, f"{not_sep}*", not_sep))
742 if idx < last_part_idx:
743 results.append(any_sep)
744 res = "".join(results)
745 return rf"(?s:{res})\Z"