Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/fsspec/core.py: 16%
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 io
4import logging
5import posixpath
6import re
7from contextlib import ExitStack
8from glob import has_magic
9from pathlib import Path
11# for backwards compat, we export cache things from here too
12from fsspec.caching import ( # noqa: F401
13 BaseCache,
14 BlockCache,
15 BytesCache,
16 MMapCache,
17 ReadAheadCache,
18 caches,
19)
20from fsspec.compression import compr
21from fsspec.config import conf
22from fsspec.registry import available_protocols, filesystem, get_filesystem_class
23from fsspec.utils import (
24 _unstrip_protocol,
25 build_name_function,
26 infer_compression,
27 stringify_path,
28)
30logger = logging.getLogger("fsspec")
33class OpenFile:
34 """
35 File-like object to be used in a context
37 Can layer (buffered) text-mode and compression over any file-system, which
38 are typically binary-only.
40 These instances are safe to serialize, as the low-level file object
41 is not created until invoked using ``with``.
43 Parameters
44 ----------
45 fs: FileSystem
46 The file system to use for opening the file. Should be a subclass or duck-type
47 with ``fsspec.spec.AbstractFileSystem``
48 path: str
49 Location to open
50 mode: str like 'rb', optional
51 Mode of the opened file
52 compression: str or None, optional
53 Compression to apply
54 encoding: str or None, optional
55 The encoding to use if opened in text mode.
56 errors: str or None, optional
57 How to handle encoding errors if opened in text mode.
58 newline: None or str
59 Passed to TextIOWrapper in text mode, how to handle line endings.
60 """
62 def __init__(
63 self,
64 fs,
65 path,
66 mode="rb",
67 compression=None,
68 encoding=None,
69 errors=None,
70 newline=None,
71 ):
72 self.fs = fs
73 self.path = path
74 self.mode = mode
75 self.compression = get_compression(path, compression)
76 self.encoding = encoding
77 self.errors = errors
78 self.newline = newline
79 self.fobjects = []
81 def __reduce__(self):
82 return (
83 OpenFile,
84 (
85 self.fs,
86 self.path,
87 self.mode,
88 self.compression,
89 self.encoding,
90 self.errors,
91 self.newline,
92 ),
93 )
95 def __repr__(self):
96 return f"<OpenFile '{self.path}'>"
98 def __enter__(self):
99 mode = self.mode.replace("t", "").replace("b", "") + "b"
101 try:
102 f = self.fs.open(self.path, mode=mode)
103 except FileNotFoundError as e:
104 if has_magic(self.path):
105 raise FileNotFoundError(
106 f"{self.path} not found. "
107 "The URL contains glob characters: you maybe needed\n"
108 "to pass expand=True in fsspec.open() or the storage_options of \n"
109 "your library. You can also set the config value 'open_expand'\n"
110 "before import, or fsspec.core.DEFAULT_EXPAND at runtime, to True."
111 ) from e
112 raise
114 self.fobjects = [f]
116 try:
117 if self.compression is not None:
118 compress = compr[self.compression]
119 f = compress(f, mode=mode[0])
120 self.fobjects.append(f)
122 if "b" not in self.mode:
123 # assume, for example, that 'r' is equivalent to 'rt' as in builtin
124 f = PickleableTextIOWrapper(
125 f, encoding=self.encoding, errors=self.errors, newline=self.newline
126 )
127 self.fobjects.append(f)
128 except BaseException:
129 self.close()
130 raise
132 return self.fobjects[-1]
134 def __exit__(self, *args):
135 self.close()
137 @property
138 def full_name(self):
139 return _unstrip_protocol(self.path, self.fs)
141 def open(self):
142 """Materialise this as a real open file without context
144 The OpenFile object should be explicitly closed to avoid enclosed file
145 instances persisting. You must, therefore, keep a reference to the OpenFile
146 during the life of the file-like it generates.
147 """
148 return self.__enter__()
150 def close(self):
151 """Close all encapsulated file objects"""
153 def close_one(f):
154 try:
155 if "r" not in self.mode and not f.closed:
156 f.flush()
157 finally:
158 f.close()
160 # Unwind in reverse order, attempting every close even if one fails.
161 with ExitStack() as stack:
162 for f in self.fobjects:
163 stack.callback(close_one, f)
164 self.fobjects.clear()
167class OpenFiles(list):
168 """List of OpenFile instances
170 Can be used in a single context, which opens and closes all of the
171 contained files. Normal list access to get the elements works as
172 normal.
174 A special case is made for caching filesystems - the files will
175 be down/uploaded together at the start or end of the context, and
176 this may happen concurrently, if the target filesystem supports it.
177 """
179 def __init__(self, *args, mode="rb", fs=None):
180 self.mode = mode
181 self.fs = fs
182 self.files = []
183 super().__init__(*args)
185 def __enter__(self):
186 if self.fs is None:
187 raise ValueError("Context has already been used")
189 fs = self.fs
190 while True:
191 if hasattr(fs, "open_many"):
192 # check for concurrent cache download; or set up for upload
193 self.files = fs.open_many(self)
194 return self.files
195 if hasattr(fs, "fs") and fs.fs is not None:
196 fs = fs.fs
197 else:
198 break
199 with ExitStack() as stack:
200 files = [stack.enter_context(s) for s in self]
201 # All files opened: transfer cleanup to OpenFiles.__exit__ so they
202 # stay open after this temporary entry-failure guard exits.
203 stack.pop_all()
204 return files
206 def __exit__(self, *args):
207 fs = self.fs
208 with ExitStack() as stack:
209 if "r" in self.mode:
210 # open_many() returns files without populating the OpenFile objects.
211 for f in reversed(self.files):
212 stack.callback(f.close)
213 for s in reversed(self):
214 stack.callback(s.close)
215 if "r" not in self.mode:
216 while True:
217 if hasattr(fs, "open_many"):
218 # check for concurrent cache upload
219 fs.commit_many(self.files)
220 return
221 if hasattr(fs, "fs") and fs.fs is not None:
222 fs = fs.fs
223 else:
224 break
226 def __getitem__(self, item):
227 out = super().__getitem__(item)
228 if isinstance(item, slice):
229 return OpenFiles(out, mode=self.mode, fs=self.fs)
230 return out
232 def __repr__(self):
233 return f"<List of {len(self)} OpenFile instances>"
236def open_files(
237 urlpath,
238 mode="rb",
239 compression=None,
240 encoding="utf8",
241 errors=None,
242 name_function=None,
243 num=1,
244 protocol=None,
245 newline=None,
246 auto_mkdir=True,
247 expand=True,
248 **kwargs,
249):
250 """Given a path or paths, return a list of ``OpenFile`` objects.
252 For writing, a str path must contain the "*" character, which will be filled
253 in by increasing numbers, e.g., "part*" -> "part1", "part2" if num=2.
255 For either reading or writing, can instead provide explicit list of paths.
257 Parameters
258 ----------
259 urlpath: string or list
260 Absolute or relative filepath(s). Prefix with a protocol like ``s3://``
261 to read from alternative filesystems. To read from multiple files you
262 can pass a globstring or a list of paths, with the caveat that they
263 must all have the same protocol.
264 mode: 'rb', 'wt', etc.
265 compression: string or None
266 If given, open file using compression codec. Can either be a compression
267 name (a key in ``fsspec.compression.compr``) or "infer" to guess the
268 compression from the filename suffix.
269 encoding: str
270 For text mode only
271 errors: None or str
272 Passed to TextIOWrapper in text mode
273 name_function: function or None
274 if opening a set of files for writing, those files do not yet exist,
275 so we need to generate their names by formatting the urlpath for
276 each sequence number
277 num: int [1]
278 if writing mode, number of files we expect to create (passed to
279 name+function)
280 protocol: str or None
281 If given, overrides the protocol found in the URL.
282 newline: bytes or None
283 Used for line terminator in text mode. If None, uses system default;
284 if blank, uses no translation.
285 auto_mkdir: bool (True)
286 If in write mode, this will ensure the target directory exists before
287 writing, by calling ``fs.mkdirs(exist_ok=True)``.
288 expand: bool
289 **kwargs: dict
290 Extra options that make sense to a particular storage connection, e.g.
291 host, port, username, password, etc.
293 Examples
294 --------
295 >>> files = open_files('2015-*-*.csv') # doctest: +SKIP
296 >>> files = open_files(
297 ... 's3://bucket/2015-*-*.csv.gz', compression='gzip'
298 ... ) # doctest: +SKIP
300 Returns
301 -------
302 An ``OpenFiles`` instance, which is a list of ``OpenFile`` objects that can
303 be used as a single context
305 Notes
306 -----
307 For a full list of the available protocols and the implementations that
308 they map across to see the latest online documentation:
310 - For implementations built into ``fsspec`` see
311 https://filesystem-spec.readthedocs.io/en/latest/api.html#built-in-implementations
312 - For implementations in separate packages see
313 https://filesystem-spec.readthedocs.io/en/latest/api.html#other-known-implementations
314 """
315 fs, fs_token, paths = get_fs_token_paths(
316 urlpath,
317 mode,
318 num=num,
319 name_function=name_function,
320 storage_options=kwargs,
321 protocol=protocol,
322 expand=expand,
323 )
324 if fs.protocol == "file":
325 fs.auto_mkdir = auto_mkdir
326 elif "r" not in mode and auto_mkdir:
327 parents = {fs._parent(path) for path in paths}
328 for parent in parents:
329 try:
330 fs.makedirs(parent, exist_ok=True)
331 except PermissionError:
332 pass
333 return OpenFiles(
334 [
335 OpenFile(
336 fs,
337 path,
338 mode=mode,
339 compression=compression,
340 encoding=encoding,
341 errors=errors,
342 newline=newline,
343 )
344 for path in paths
345 ],
346 mode=mode,
347 fs=fs,
348 )
351def _un_chain(path, kwargs):
352 # Avoid a circular import
353 from fsspec.implementations.chained import ChainedFileSystem
355 if "::" in path:
356 x = re.compile(".*[^a-z]+.*") # test for non protocol-like single word
357 known_protocols = set(available_protocols())
358 bits = []
360 # split on '::', then ensure each bit has a protocol
361 for p in path.split("::"):
362 if p in known_protocols:
363 bits.append(p + "://")
364 elif "://" in p or x.match(p):
365 bits.append(p)
366 else:
367 bits.append(p + "://")
368 else:
369 bits = [path]
371 # [[url, protocol, kwargs], ...]
372 out = []
373 previous_bit = None
374 kwargs = kwargs.copy()
376 for bit in reversed(bits):
377 protocol = kwargs.pop("protocol", None) or split_protocol(bit)[0] or "file"
378 cls = get_filesystem_class(protocol)
379 extra_kwargs = cls._get_kwargs_from_urls(bit)
380 kws = kwargs.pop(protocol, {}).copy()
382 if bit is bits[0]:
383 kws.update(kwargs)
385 kw = dict(
386 **{k: v for k, v in extra_kwargs.items() if k not in kws or v != kws[k]},
387 **kws,
388 )
389 bit = cls._strip_protocol(bit)
391 if (
392 "target_protocol" not in kw
393 and issubclass(cls, ChainedFileSystem)
394 and not bit
395 ):
396 # replace bit if we are chaining and no path given
397 bit = previous_bit
399 out.append((bit, protocol, kw))
400 previous_bit = bit
402 out.reverse()
403 return out
406def url_to_fs(url, **kwargs):
407 """
408 Turn fully-qualified and potentially chained URL into filesystem instance
410 Parameters
411 ----------
412 url : str
413 The fsspec-compatible URL
414 **kwargs: dict
415 Extra options that make sense to a particular storage connection, e.g.
416 host, port, username, password, etc.
418 Returns
419 -------
420 filesystem : FileSystem
421 The new filesystem discovered from ``url`` and created with
422 ``**kwargs``.
423 urlpath : str
424 The file-systems-specific URL for ``url``.
425 """
426 url = stringify_path(url)
427 # non-FS arguments that appear in fsspec.open()
428 # inspect could keep this in sync with open()'s signature
429 known_kwargs = {
430 "compression",
431 "encoding",
432 "errors",
433 "expand",
434 "mode",
435 "name_function",
436 "newline",
437 "num",
438 }
439 kwargs = {k: v for k, v in kwargs.items() if k not in known_kwargs}
440 chain = _un_chain(url, kwargs)
441 inkwargs = {}
442 # Reverse iterate the chain, creating a nested target_* structure
443 for i, ch in enumerate(reversed(chain)):
444 urls, protocol, kw = ch
445 if i == len(chain) - 1:
446 inkwargs = dict(**kw, **inkwargs)
447 continue
448 inkwargs["target_options"] = dict(**kw, **inkwargs)
449 inkwargs["target_protocol"] = protocol
450 inkwargs["fo"] = urls
451 urlpath, protocol, _ = chain[0]
452 fs = filesystem(protocol, **inkwargs)
453 return fs, urlpath
456DEFAULT_EXPAND = conf.get("open_expand", False)
459def open(
460 urlpath,
461 mode="rb",
462 compression=None,
463 encoding="utf8",
464 errors=None,
465 protocol=None,
466 newline=None,
467 expand=None,
468 **kwargs,
469) -> OpenFile:
470 """Given a path or paths, return one ``OpenFile`` object.
472 Parameters
473 ----------
474 urlpath: string or list
475 Absolute or relative filepath. Prefix with a protocol like ``s3://``
476 to read from alternative filesystems. Should not include glob
477 character(s).
478 mode: 'rb', 'wt', etc.
479 compression: string or None
480 If given, open file using compression codec. Can either be a compression
481 name (a key in ``fsspec.compression.compr``) or "infer" to guess the
482 compression from the filename suffix.
483 encoding: str
484 For text mode only
485 errors: None or str
486 Passed to TextIOWrapper in text mode
487 protocol: str or None
488 If given, overrides the protocol found in the URL.
489 newline: bytes or None
490 Used for line terminator in text mode. If None, uses system default;
491 if blank, uses no translation.
492 expand: bool or None
493 Whether to regard file paths containing special glob characters as needing
494 expansion (finding the first match) or absolute. Setting False allows using
495 paths which do embed such characters. If None (default), this argument
496 takes its value from the DEFAULT_EXPAND module variable, which takes
497 its initial value from the "open_expand" config value at startup, which will
498 be False if not set.
499 **kwargs: dict
500 Extra options that make sense to a particular storage connection, e.g.
501 host, port, username, password, etc.
503 Examples
504 --------
505 >>> openfile = open('2015-01-01.csv') # doctest: +SKIP
506 >>> openfile = open(
507 ... 's3://bucket/2015-01-01.csv.gz', compression='gzip'
508 ... ) # doctest: +SKIP
509 >>> with openfile as f:
510 ... df = pd.read_csv(f) # doctest: +SKIP
511 ...
513 Returns
514 -------
515 ``OpenFile`` object.
517 Notes
518 -----
519 For a full list of the available protocols and the implementations that
520 they map across to see the latest online documentation:
522 - For implementations built into ``fsspec`` see
523 https://filesystem-spec.readthedocs.io/en/latest/api.html#built-in-implementations
524 - For implementations in separate packages see
525 https://filesystem-spec.readthedocs.io/en/latest/api.html#other-known-implementations
526 """
527 expand = DEFAULT_EXPAND if expand is None else expand
528 out = open_files(
529 urlpath=[urlpath],
530 mode=mode,
531 compression=compression,
532 encoding=encoding,
533 errors=errors,
534 protocol=protocol,
535 newline=newline,
536 expand=expand,
537 **kwargs,
538 )
539 if not out:
540 raise FileNotFoundError(urlpath)
541 return out[0]
544def open_local(
545 url: str | list[str] | Path | list[Path],
546 mode: str = "rb",
547 **storage_options: dict,
548) -> str | list[str]:
549 """Open file(s) which can be resolved to local
551 For files which either are local, or get downloaded upon open
552 (e.g., by file caching)
554 Parameters
555 ----------
556 url: str or list(str)
557 mode: str
558 Must be read mode
559 storage_options:
560 passed on to FS for or used by open_files (e.g., compression)
561 """
562 if "r" not in mode:
563 raise ValueError("Can only ensure local files when reading")
564 of = open_files(url, mode=mode, **storage_options)
565 if not getattr(of[0].fs, "local_file", False):
566 raise ValueError(
567 "open_local can only be used on a filesystem which"
568 " has attribute local_file=True"
569 )
570 with of as files:
571 paths = [f.name for f in files]
572 if (isinstance(url, str) and not has_magic(url)) or isinstance(url, Path):
573 return paths[0]
574 return paths
577def get_compression(urlpath, compression):
578 if compression == "infer":
579 compression = infer_compression(urlpath)
580 if compression is not None and compression not in compr:
581 raise ValueError(f"Compression type {compression} not supported")
582 return compression
585def split_protocol(urlpath):
586 """Return protocol, path pair"""
587 urlpath = stringify_path(urlpath)
588 if "://" in urlpath:
589 protocol, path = urlpath.split("://", 1)
590 if len(protocol) > 1:
591 # excludes Windows paths
592 return protocol, path
593 if urlpath.startswith("data:"):
594 return urlpath.split(":", 1)
595 return None, urlpath
598def strip_protocol(urlpath):
599 """Return only path part of full URL, according to appropriate backend"""
600 protocol, _ = split_protocol(urlpath)
601 cls = get_filesystem_class(protocol)
602 return cls._strip_protocol(urlpath)
605def expand_paths_if_needed(paths, mode, num, fs, name_function):
606 """Expand paths if they have a ``*`` in them (write mode) or any of ``*?[]``
607 in them (read mode).
609 :param paths: list of paths
610 mode: str
611 Mode in which to open files.
612 num: int
613 If opening in writing mode, number of files we expect to create.
614 fs: filesystem object
615 name_function: callable
616 If opening in writing mode, this callable is used to generate path
617 names. Names are generated for each partition by
618 ``urlpath.replace('*', name_function(partition_index))``.
619 :return: list of paths
620 """
621 expanded_paths = []
622 paths = list(paths)
624 if "w" in mode or "x" in mode: # write mode
625 if sum(1 for p in paths if "*" in p) > 1:
626 raise ValueError(
627 "When writing data, only one filename mask can be specified."
628 )
629 num = max(num, len(paths))
631 for curr_path in paths:
632 if "*" in curr_path:
633 # expand using name_function
634 expanded_paths.extend(_expand_paths(curr_path, name_function, num))
635 else:
636 expanded_paths.append(curr_path)
637 # if we generated more paths that asked for, trim the list
638 if len(expanded_paths) > num:
639 expanded_paths = expanded_paths[:num]
641 else: # read mode
642 for curr_path in paths:
643 if has_magic(curr_path):
644 # expand using glob
645 expanded_paths.extend(fs.glob(curr_path))
646 else:
647 expanded_paths.append(curr_path)
649 return expanded_paths
652def get_fs_token_paths(
653 urlpath,
654 mode="rb",
655 num=1,
656 name_function=None,
657 storage_options=None,
658 protocol=None,
659 expand=True,
660):
661 """Filesystem, deterministic token, and paths from a urlpath and options.
663 Parameters
664 ----------
665 urlpath: string or iterable
666 Absolute or relative filepath, URL (may include protocols like
667 ``s3://``), or globstring pointing to data.
668 mode: str, optional
669 Mode in which to open files.
670 num: int, optional
671 If opening in writing mode, number of files we expect to create.
672 name_function: callable, optional
673 If opening in writing mode, this callable is used to generate path
674 names. Names are generated for each partition by
675 ``urlpath.replace('*', name_function(partition_index))``.
676 storage_options: dict, optional
677 Additional keywords to pass to the filesystem class.
678 protocol: str or None
679 To override the protocol specifier in the URL
680 expand: bool
681 Expand string paths for writing, assuming the path is a directory
682 """
683 if isinstance(urlpath, (list, tuple, set)):
684 if not urlpath:
685 raise ValueError("empty urlpath sequence")
686 urlpath0 = stringify_path(next(iter(urlpath)))
687 else:
688 urlpath0 = stringify_path(urlpath)
689 storage_options = storage_options or {}
690 if protocol:
691 storage_options["protocol"] = protocol
692 chain = _un_chain(urlpath0, storage_options or {})
693 inkwargs = {}
694 # Reverse iterate the chain, creating a nested target_* structure
695 for i, ch in enumerate(reversed(chain)):
696 urls, nested_protocol, kw = ch
697 if i == len(chain) - 1:
698 inkwargs = dict(**kw, **inkwargs)
699 continue
700 inkwargs["target_options"] = dict(**kw, **inkwargs)
701 inkwargs["target_protocol"] = nested_protocol
702 inkwargs["fo"] = urls
703 paths, protocol, _ = chain[0]
704 fs = filesystem(protocol, **inkwargs)
705 if isinstance(urlpath, (list, tuple, set)):
706 pchains = [
707 _un_chain(stringify_path(u), storage_options or {})[0] for u in urlpath
708 ]
709 if len({pc[1] for pc in pchains}) > 1:
710 raise ValueError(f"Protocol mismatch getting fs from {urlpath}")
711 paths = [pc[0] for pc in pchains]
712 else:
713 paths = fs._strip_protocol(paths)
714 if isinstance(paths, (list, tuple, set)):
715 if expand:
716 paths = expand_paths_if_needed(paths, mode, num, fs, name_function)
717 elif not isinstance(paths, list):
718 paths = list(paths)
719 else:
720 if ("w" in mode or "x" in mode) and expand:
721 paths = _expand_paths(paths, name_function, num)
722 elif "*" in paths:
723 paths = [f for f in sorted(fs.glob(paths)) if not fs.isdir(f)]
724 else:
725 paths = [paths]
727 return fs, fs._fs_token, paths
730def _expand_paths(path, name_function, num):
731 if isinstance(path, str):
732 if path.count("*") > 1:
733 raise ValueError("Output path spec must contain exactly one '*'.")
734 elif "*" not in path:
735 # paths are "/"-separated on every filesystem, including local ones
736 path = posixpath.join(path, "*.part")
738 if name_function is None:
739 name_function = build_name_function(num - 1)
741 paths = [path.replace("*", name_function(i)) for i in range(num)]
742 if paths != sorted(paths):
743 logger.warning(
744 "In order to preserve order between partitions"
745 " paths created with ``name_function`` should "
746 "sort to partition order"
747 )
748 elif isinstance(path, (tuple, list)):
749 assert len(path) == num
750 paths = list(path)
751 else:
752 raise ValueError(
753 "Path should be either\n"
754 "1. A list of paths: ['foo.json', 'bar.json', ...]\n"
755 "2. A directory: 'foo/\n"
756 "3. A path with a '*' in it: 'foo.*.json'"
757 )
758 return paths
761class PickleableTextIOWrapper(io.TextIOWrapper):
762 """TextIOWrapper cannot be pickled. This solves it.
764 Requires that ``buffer`` be pickleable, which all instances of
765 AbstractBufferedFile are.
766 """
768 def __init__(
769 self,
770 buffer,
771 encoding=None,
772 errors=None,
773 newline=None,
774 line_buffering=False,
775 write_through=False,
776 ):
777 self.args = buffer, encoding, errors, newline, line_buffering, write_through
778 super().__init__(*self.args)
780 def __reduce__(self):
781 return PickleableTextIOWrapper, self.args