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 Callable, Iterable, Iterator, Sequence
11from functools import partial
12from hashlib import md5
13from importlib.metadata import version
14from typing import IO, TYPE_CHECKING, Any, TypeVar
15from urllib.parse import urlsplit
17if TYPE_CHECKING:
18 import pathlib
19 from typing import TypeGuard
21 from fsspec.spec import AbstractFileSystem
24DEFAULT_BLOCK_SIZE = 5 * 2**20
26T = TypeVar("T")
29def infer_storage_options(
30 urlpath: str, inherit_storage_options: dict[str, Any] | None = None
31) -> dict[str, Any]:
32 """Infer storage options from URL path and merge it with existing storage
33 options.
35 Parameters
36 ----------
37 urlpath: str or unicode
38 Either local absolute file path or URL (hdfs://namenode:8020/file.csv)
39 inherit_storage_options: dict (optional)
40 Its contents will get merged with the inferred information from the
41 given path
43 Returns
44 -------
45 Storage options dict.
47 Examples
48 --------
49 >>> infer_storage_options('/mnt/datasets/test.csv') # doctest: +SKIP
50 {"protocol": "file", "path", "/mnt/datasets/test.csv"}
51 >>> infer_storage_options(
52 ... 'hdfs://username:pwd@node:123/mnt/datasets/test.csv?q=1',
53 ... inherit_storage_options={'extra': 'value'},
54 ... ) # doctest: +SKIP
55 {"protocol": "hdfs", "username": "username", "password": "pwd",
56 "host": "node", "port": 123, "path": "/mnt/datasets/test.csv",
57 "url_query": "q=1", "extra": "value"}
58 """
60 # Discover Windows paths including disk name in this special case.
61 is_filesystem = re.match(r"^[a-zA-Z]:[\\/]", urlpath)
63 # Discover URI according to RFC 3986: Scheme names consist of a
64 # sequence of characters beginning with a letter and followed by
65 # any combination of letters, digits, plus ("+"), period ("."),
66 # or hyphen ("-").
67 # https://datatracker.ietf.org/doc/html/rfc3986#section-3.1
68 is_uri = re.match(r"^[a-zA-Z0-9+.-]+://", urlpath)
70 if is_filesystem or is_uri is None:
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 # Only consider the final path component: a "." in a parent directory name
444 # (e.g. "/path/to.dir/file") is not the file's extension.
445 ext_parts = url.rsplit("/", 1)[-1].rsplit(".", 1)
446 if len(ext_parts) > 1:
447 return ext_parts[-1]
448 return ""
451def can_be_local(path: str) -> bool:
452 """Can the given URL be used with open_local?"""
453 from fsspec import get_filesystem_class
455 try:
456 return getattr(get_filesystem_class(get_protocol(path)), "local_file", False)
457 except (ValueError, ImportError):
458 # not in registry or import failed
459 return False
462def get_package_version_without_import(name: str) -> str | None:
463 """For given package name, try to find the version without importing it
465 Import and package.__version__ is still the backup here, so an import
466 *might* happen.
468 Returns either the version string, or None if the package
469 or the version was not readily found.
470 """
471 if name in sys.modules:
472 mod = sys.modules[name]
473 if hasattr(mod, "__version__"):
474 return mod.__version__
475 try:
476 return version(name)
477 except: # noqa: E722
478 pass
479 try:
480 import importlib
482 mod = importlib.import_module(name)
483 return mod.__version__
484 except (ImportError, AttributeError):
485 return None
488def setup_logging(
489 logger: logging.Logger | None = None,
490 logger_name: str | None = None,
491 level: str = "DEBUG",
492 clear: bool = True,
493) -> logging.Logger:
494 if logger is None and logger_name is None:
495 raise ValueError("Provide either logger object or logger name")
496 logger = logger or logging.getLogger(logger_name)
497 handle = logging.StreamHandler()
498 formatter = logging.Formatter(
499 "%(asctime)s - %(name)s - %(levelname)s - %(funcName)s -- %(message)s"
500 )
501 handle.setFormatter(formatter)
502 if clear:
503 logger.handlers.clear()
504 logger.addHandler(handle)
505 logger.setLevel(level)
506 return logger
509def _unstrip_protocol(name: str, fs: AbstractFileSystem) -> str:
510 return fs.unstrip_protocol(name)
513def mirror_from(
514 origin_name: str, methods: Iterable[str]
515) -> Callable[[type[T]], type[T]]:
516 """Mirror attributes and methods from the given
517 origin_name attribute of the instance to the
518 decorated class"""
520 def origin_getter(method: str, self: Any) -> Any:
521 origin = getattr(self, origin_name)
522 return getattr(origin, method)
524 def wrapper(cls: type[T]) -> type[T]:
525 for method in methods:
526 wrapped_method = partial(origin_getter, method)
527 setattr(cls, method, property(wrapped_method))
528 return cls
530 return wrapper
533@contextlib.contextmanager
534def nullcontext(obj: T) -> Iterator[T]:
535 yield obj
538def merge_offset_ranges(
539 paths: list[str],
540 starts: list[int] | int,
541 ends: list[int] | int,
542 max_gap: int = 0,
543 max_block: int | None = None,
544 sort: bool = True,
545) -> tuple[list[str], list[int], list[int]]:
546 """Merge adjacent byte-offset ranges when the inter-range
547 gap is <= `max_gap`, and when the merged byte range does not
548 exceed `max_block` (if specified). By default, this function
549 will re-order the input paths and byte ranges to ensure sorted
550 order. If the user can guarantee that the inputs are already
551 sorted, passing `sort=False` will skip the re-ordering.
552 """
553 # Check input
554 if not isinstance(paths, list):
555 raise TypeError
556 if not isinstance(starts, list):
557 starts = [starts] * len(paths)
558 if not isinstance(ends, list):
559 ends = [ends] * len(paths)
560 if len(starts) != len(paths) or len(ends) != len(paths):
561 raise ValueError
563 # Early Return
564 if len(starts) <= 1:
565 return paths, starts, ends
567 starts = [s or 0 for s in starts]
568 # Sort by paths and then ranges if `sort=True`
569 if sort:
570 paths, starts, ends = (
571 list(v)
572 for v in zip(
573 *sorted(
574 zip(paths, starts, ends),
575 )
576 )
577 )
578 remove = []
579 for i, (path, start, end) in enumerate(zip(paths, starts, ends)):
580 if any(
581 e is not None and p == path and start >= s and end <= e and i != i2
582 for i2, (p, s, e) in enumerate(zip(paths, starts, ends))
583 ):
584 remove.append(i)
585 paths = [p for i, p in enumerate(paths) if i not in remove]
586 starts = [s for i, s in enumerate(starts) if i not in remove]
587 ends = [e for i, e in enumerate(ends) if i not in remove]
589 if paths:
590 # Loop through the coupled `paths`, `starts`, and
591 # `ends`, and merge adjacent blocks when appropriate
592 new_paths = paths[:1]
593 new_starts = starts[:1]
594 new_ends = ends[:1]
595 for i in range(1, len(paths)):
596 if paths[i] == paths[i - 1] and new_ends[-1] is None:
597 continue
598 elif (
599 paths[i] != paths[i - 1]
600 or ((starts[i] - new_ends[-1]) > max_gap)
601 or (max_block is not None and (ends[i] - new_starts[-1]) > max_block)
602 ):
603 # Cannot merge with previous block.
604 # Add new `paths`, `starts`, and `ends` elements
605 new_paths.append(paths[i])
606 new_starts.append(starts[i])
607 new_ends.append(ends[i])
608 else:
609 # Merge with the previous block by updating the
610 # last element of `ends`
611 new_ends[-1] = ends[i]
612 return new_paths, new_starts, new_ends
614 # `paths` is empty. Just return input lists
615 return paths, starts, ends
618def file_size(filelike: IO[bytes]) -> int:
619 """Find length of any open read-mode file-like"""
620 pos = filelike.tell()
621 try:
622 return filelike.seek(0, 2)
623 finally:
624 filelike.seek(pos)
627@contextlib.contextmanager
628def atomic_write(path: str, mode: str = "wb"):
629 """
630 A context manager that opens a temporary file next to `path` and, on exit,
631 replaces `path` with the temporary file, thereby updating `path`
632 atomically.
633 """
634 fd, fn = tempfile.mkstemp(
635 dir=os.path.dirname(path), prefix=os.path.basename(path) + "-"
636 )
637 try:
638 with open(fd, mode) as fp:
639 yield fp
640 except BaseException:
641 with contextlib.suppress(FileNotFoundError):
642 os.unlink(fn)
643 raise
644 else:
645 os.replace(fn, path)
648def _translate(pat, STAR, QUESTION_MARK):
649 # Copied from: https://github.com/python/cpython/pull/106703.
650 res: list[str] = []
651 add = res.append
652 i, n = 0, len(pat)
653 while i < n:
654 c = pat[i]
655 i = i + 1
656 if c == "*":
657 # compress consecutive `*` into one
658 if (not res) or res[-1] is not STAR:
659 add(STAR)
660 elif c == "?":
661 add(QUESTION_MARK)
662 elif c == "[":
663 j = i
664 if j < n and pat[j] == "!":
665 j = j + 1
666 if j < n and pat[j] == "]":
667 j = j + 1
668 while j < n and pat[j] != "]":
669 j = j + 1
670 if j >= n:
671 add("\\[")
672 else:
673 stuff = pat[i:j]
674 if "-" not in stuff:
675 stuff = stuff.replace("\\", r"\\")
676 else:
677 chunks = []
678 k = i + 2 if pat[i] == "!" else i + 1
679 while True:
680 k = pat.find("-", k, j)
681 if k < 0:
682 break
683 chunks.append(pat[i:k])
684 i = k + 1
685 k = k + 3
686 chunk = pat[i:j]
687 if chunk:
688 chunks.append(chunk)
689 else:
690 chunks[-1] += "-"
691 # Remove empty ranges -- invalid in RE.
692 for k in range(len(chunks) - 1, 0, -1):
693 if chunks[k - 1][-1] > chunks[k][0]:
694 chunks[k - 1] = chunks[k - 1][:-1] + chunks[k][1:]
695 del chunks[k]
696 # Escape backslashes and hyphens for set difference (--).
697 # Hyphens that create ranges shouldn't be escaped.
698 stuff = "-".join(
699 s.replace("\\", r"\\").replace("-", r"\-") for s in chunks
700 )
701 # Escape set operations (&&, ~~ and ||).
702 stuff = re.sub(r"([&~|])", r"\\\1", stuff)
703 i = j + 1
704 if not stuff:
705 # Empty range: never match.
706 add("(?!)")
707 elif stuff == "!":
708 # Negated empty range: match any character.
709 add(".")
710 else:
711 if stuff[0] == "!":
712 stuff = "^" + stuff[1:]
713 elif stuff[0] in ("^", "["):
714 stuff = "\\" + stuff
715 add(f"[{stuff}]")
716 else:
717 add(re.escape(c))
718 assert i == n
719 return res
722def glob_translate(pat):
723 # Copied from: https://github.com/python/cpython/pull/106703.
724 # The keyword parameters' values are fixed to:
725 # recursive=True, include_hidden=True, seps=None
726 """Translate a pathname with shell wildcards to a regular expression."""
727 if os.path.altsep:
728 seps = os.path.sep + os.path.altsep
729 else:
730 seps = os.path.sep
731 escaped_seps = "".join(map(re.escape, seps))
732 any_sep = f"[{escaped_seps}]" if len(seps) > 1 else escaped_seps
733 not_sep = f"[^{escaped_seps}]"
734 one_last_segment = f"{not_sep}+"
735 one_segment = f"{one_last_segment}{any_sep}"
736 any_segments = f"(?:.+{any_sep})?"
737 any_last_segments = ".*"
738 results = []
739 parts = re.split(any_sep, pat)
740 last_part_idx = len(parts) - 1
741 for idx, part in enumerate(parts):
742 if part == "*":
743 results.append(one_segment if idx < last_part_idx else one_last_segment)
744 continue
745 if part == "**":
746 results.append(any_segments if idx < last_part_idx else any_last_segments)
747 continue
748 elif "**" in part:
749 raise ValueError(
750 "Invalid pattern: '**' can only be an entire path component"
751 )
752 if part:
753 results.extend(_translate(part, f"{not_sep}*", not_sep))
754 if idx < last_part_idx:
755 results.append(any_sep)
756 res = "".join(results)
757 return rf"(?s:{res})\Z"