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
1#
2# Copyright (C) 2019 Radim Rehurek <me@radimrehurek.com>
3#
4# This code is distributed under the terms and conditions
5# from the MIT License (MIT).
6#
8"""Implements the majority of smart_open's top-level API."""
10from __future__ import annotations
12import collections
13import contextlib
14import locale
15import logging
16import os
17import os.path
18import pathlib
19import urllib.parse
20from typing import IO, TYPE_CHECKING, Any, BinaryIO, Literal, TextIO, cast, overload
22import smart_open.compression as so_compression
24#
25# This module defines a function called smart_open so we cannot use
26# smart_open.submodule to reference to the submodules.
27#
28import smart_open.local_file as so_file
29import smart_open.utils as so_utils
30from smart_open import doctools, transport
32if TYPE_CHECKING:
33 from collections.abc import Callable
35 from typing_extensions import Self
37 from smart_open._typing import CompressionKwargs, TransportParams, Uri
39logger = logging.getLogger(__name__)
41DEFAULT_ENCODING = locale.getpreferredencoding(do_setlocale=False)
44def _sniff_scheme(uri_as_string: str) -> str:
45 """Returns the scheme of the URL only, as a string."""
46 #
47 # urlsplit doesn't work on Windows -- it parses the drive as the scheme...
48 # no protocol given => assume a local file
49 #
50 if os.name == "nt" and "://" not in uri_as_string:
51 uri_as_string = "file://" + uri_as_string
53 return urllib.parse.urlsplit(uri_as_string).scheme
56def parse_uri(uri_as_string: str) -> tuple[Any, ...]:
57 """Parse the given URI from a string.
59 Args:
60 uri_as_string: The URI to parse.
62 Returns:
63 The parsed URI as a ``collections.namedtuple``.
65 smart_open/doctools.py magic goes here
66 """
67 scheme = _sniff_scheme(uri_as_string)
68 submodule = transport.get_transport(scheme)
69 as_dict = submodule.parse_uri(uri_as_string)
71 #
72 # The conversion to a namedtuple is just to keep the old tests happy while
73 # I'm still refactoring.
74 #
75 Uri = collections.namedtuple("Uri", sorted(as_dict.keys())) # noqa: PYI024 # legacy public type
76 return Uri(**as_dict)
79#
80# To keep old unit tests happy while I'm refactoring.
81#
82_parse_uri = parse_uri
84_builtin_open = open
87@overload
88def open(
89 uri: Uri,
90 mode: Literal["r", "w", "a", "r+", "w+", "a+", "rt", "wt", "at"] = ...,
91 buffering: int = ...,
92 encoding: str | None = ...,
93 errors: str | None = ...,
94 newline: str | None = ...,
95 closefd: bool = ..., # noqa: FBT001 # public API
96 opener: Callable[[str, int], int] | None = ...,
97 compression: str = ...,
98 compression_kwargs: CompressionKwargs | None = ...,
99 transport_params: TransportParams | None = ...,
100) -> TextIO: ...
103@overload
104def open(
105 uri: Uri,
106 mode: Literal["rb", "wb", "ab", "rb+", "wb+", "ab+", "br", "bw", "ba"],
107 buffering: int = ...,
108 *,
109 encoding: None = ...,
110 errors: str | None = ...,
111 newline: str | None = ...,
112 closefd: bool = ...,
113 opener: Callable[[str, int], int] | None = ...,
114 compression: str = ...,
115 compression_kwargs: CompressionKwargs | None = ...,
116 transport_params: TransportParams | None = ...,
117) -> BinaryIO: ...
120@overload
121def open(
122 uri: Uri,
123 mode: str = ...,
124 buffering: int = ...,
125 encoding: str | None = ...,
126 errors: str | None = ...,
127 newline: str | None = ...,
128 closefd: bool = ..., # noqa: FBT001 # public API
129 opener: Callable[[str, int], int] | None = ...,
130 compression: str = ...,
131 compression_kwargs: CompressionKwargs | None = ...,
132 transport_params: TransportParams | None = ...,
133) -> IO[Any]: ...
136def open( # noqa: C901, PLR0913 # legacy public API; refactor in a dedicated PR
137 uri: Uri,
138 mode: str = "r",
139 buffering: int = -1,
140 encoding: str | None = None,
141 errors: str | None = None,
142 newline: str | None = None,
143 closefd: bool = True, # noqa: ARG001, FBT001, FBT002 # public API
144 opener: Callable[[str, int], int] | None = None, # noqa: ARG001 # public API
145 compression: str = so_compression.INFER_FROM_EXTENSION,
146 compression_kwargs: CompressionKwargs | None = None,
147 transport_params: TransportParams | None = None,
148) -> IO[Any]:
149 r"""Open the URI object, returning a file-like object.
151 The URI is usually a string in a variety of formats.
152 For a full list of examples, see the :func:`parse_uri` function.
154 The URI may also be one of:
156 - an instance of the pathlib.Path class
157 - a stream (anything that implements io.IOBase-like functionality)
159 Args:
160 uri: The object to open.
161 mode: Mimics built-in open parameter of the same name.
162 buffering: Mimics built-in open parameter of the same name.
163 encoding: Mimics built-in open parameter of the same name.
164 errors: Mimics built-in open parameter of the same name.
165 newline: Mimics built-in open parameter of the same name.
166 closefd: Mimics built-in open parameter of the same name. Ignored.
167 opener: Mimics built-in open parameter of the same name. Ignored.
168 compression: Explicitly specify the compression/decompression behavior.
169 See ``smart_open.compression.get_supported_compression_types``.
170 compression_kwargs: Keyword arguments forwarded to the registered
171 compressor callback. When omitted, each library's own default level
172 applies: .gz and .bz2 default to 9 (already their maximum), while
173 .xz defaults to 6 (max 9), .zst to 3 (max 22), and .lz4 to 0 (max
174 16). To request maximum compression, pass ``{'compresslevel': 9}``
175 for .gz/.bz2, ``{'preset': 9}`` for .xz, ``{'level': 22}`` for .zst,
176 or ``{'compression_level': 16}`` for .lz4. Ignored when compression
177 is 'disable' or the URI's extension doesn't match a registered
178 compressor.
179 transport_params: Additional parameters for the transport layer (see
180 notes below).
182 Returns:
183 A file-like object.
185 Raises:
186 TypeError: If ``mode`` is not a string or if the URI type is not
187 recognized.
188 ValueError: If ``compression`` is not a supported value.
189 NotImplementedError: If ``mode`` cannot be parsed into a valid binary
190 mode.
192 Note:
193 smart_open has several implementations for its transport layer
194 (e.g. S3, HTTP). Each transport layer has a different set of keyword
195 arguments for overriding default behavior. If you specify a keyword
196 argument that is *not* supported by the transport layer being used,
197 smart_open will ignore that argument and log a warning message.
199 smart_open/doctools.py magic goes here
201 See Also:
202 - `Standard library reference <https://docs.python.org/3.14/library/functions.html#open>`__
203 - `smart_open README.md
204 <https://github.com/piskvorky/smart_open/blob/master/README.md>`__
205 """
206 if not isinstance(mode, str):
207 msg = "mode should be a string"
208 raise TypeError(msg)
210 if compression not in so_compression.get_supported_compression_types():
211 msg = f"invalid compression type: {compression}"
212 raise ValueError(msg)
214 if transport_params is None:
215 transport_params = {}
217 fobj = _shortcut_open(
218 uri,
219 mode,
220 compression=compression,
221 buffering=buffering,
222 encoding=encoding,
223 errors=errors,
224 newline=newline,
225 )
226 if fobj is not None:
227 return fobj
229 #
230 # This is a work-around for the problem described in Issue #144.
231 # If the user has explicitly specified an encoding, then assume they want
232 # us to open the destination in text mode, instead of the default binary.
233 #
234 # If we change the default mode to be text, and match the normal behavior
235 # of Py2 and 3, then the above assumption will be unnecessary.
236 #
237 if encoding is not None and "b" in mode:
238 mode = mode.replace("b", "")
240 if isinstance(uri, pathlib.Path):
241 uri = str(uri)
243 explicit_encoding = encoding
244 encoding = explicit_encoding or DEFAULT_ENCODING
246 #
247 # This is how we get from the filename to the end result. Decompression is
248 # optional, but it always accepts bytes and returns bytes.
249 #
250 # Decoding is also optional, accepts bytes and returns text. The diagram
251 # below is for reading, for writing, the flow is from right to left, but
252 # the code is identical.
253 #
254 # open as binary decompress? decode?
255 # filename ---------------> bytes -------------> bytes ---------> text
256 # binary decompressed decode
257 #
259 try:
260 binary_mode = _get_binary_mode(mode)
261 except ValueError as ve:
262 raise NotImplementedError(ve.args[0]) from ve
264 binary = _open_binary_stream(uri, binary_mode, transport_params)
265 name = getattr(binary, "name", None)
266 # prefer the stream's own name; if it's not string-like (e.g. ftp socket fileno), fall back to uri
267 filename = name if isinstance(name, str) else uri if isinstance(uri, str) else None
268 decompressed = so_compression.compression_wrapper(
269 binary,
270 binary_mode,
271 compression,
272 filename=filename,
273 compression_kwargs=compression_kwargs,
274 )
276 if "b" not in mode or explicit_encoding is not None:
277 decoded = _encoding_wrapper(
278 decompressed,
279 mode,
280 encoding=encoding,
281 errors=errors,
282 newline=newline,
283 )
284 else:
285 decoded = decompressed
287 #
288 # There are some useful methods in the binary readers, e.g. to_boto3, that get
289 # hidden by the multiple layers of wrapping we just performed. Promote
290 # them so they are visible to the user.
291 #
292 if decoded != binary:
293 promoted_attrs = ["to_boto3"]
294 for attr in promoted_attrs:
295 with contextlib.suppress(AttributeError):
296 setattr(decoded, attr, getattr(binary, attr))
298 return cast("IO[Any]", so_utils.FileLikeProxy(decoded, binary))
301def _get_binary_mode(mode_str: str) -> str: # noqa: C901 # legacy internal helper; refactor in a dedicated PR
302 #
303 # https://docs.python.org/3/library/functions.html#open
304 #
305 # The order of characters in the mode parameter appears to be unspecified.
306 # The implementation follows the examples, just to be safe.
307 #
308 mode = list(mode_str)
309 binmode = []
311 if "t" in mode and "b" in mode:
312 msg = "can't have text and binary mode at once"
313 raise ValueError(msg)
315 counts = [mode.count(x) for x in "rwa"]
316 if sum(counts) > 1:
317 msg = "must have exactly one of read/write/append mode"
318 raise ValueError(msg)
320 def transfer(char: str) -> None:
321 binmode.append(mode.pop(mode.index(char)))
323 if "a" in mode:
324 transfer("a")
325 elif "w" in mode:
326 transfer("w")
327 elif "r" in mode:
328 transfer("r")
329 else:
330 msg = "Must have exactly one of read/write/append mode and at most one plus"
331 raise ValueError(msg)
333 if "b" in mode:
334 transfer("b")
335 elif "t" in mode:
336 mode.pop(mode.index("t"))
337 binmode.append("b")
338 else:
339 binmode.append("b")
341 if "+" in mode:
342 transfer("+")
344 #
345 # There shouldn't be anything left in the mode list at this stage.
346 # If there is, then either we've missed something and the implementation
347 # of this function is broken, or the original input mode is invalid.
348 #
349 if mode:
350 msg = f"invalid mode: {mode_str!r}"
351 raise ValueError(msg)
353 return "".join(binmode)
356def _shortcut_open( # noqa: PLR0913 # legacy internal helper; refactor in a dedicated PR
357 uri: Uri,
358 mode: str,
359 compression: str,
360 buffering: int = -1,
361 encoding: str | None = None,
362 errors: str | None = None,
363 newline: str | None = None,
364) -> IO[Any] | None:
365 """Try to open the URI using the standard library io.open function.
367 This can be much faster than the alternative of opening in binary mode and
368 then decoding.
370 This is only possible under the following conditions:
372 1. Opening a local file; and
373 2. Compression is disabled
375 If it is not possible to use the built-in open for the specified URI,
376 returns None.
378 Args:
379 uri: A string indicating what to open.
380 mode: The mode to pass to the open function.
381 compression: The compression type selected.
382 buffering: Mimics built-in open parameter of the same name.
383 encoding: Mimics built-in open parameter of the same name.
384 errors: Mimics built-in open parameter of the same name.
385 newline: Mimics built-in open parameter of the same name.
387 Returns:
388 The opened file, or None if no shortcut is possible.
389 """
390 if not isinstance(uri, str):
391 return None
393 scheme = _sniff_scheme(uri)
394 if scheme not in (transport.NO_SCHEME, so_file.SCHEME):
395 return None
397 local_path = so_file.extract_local_path(uri)
398 if compression == so_compression.INFER_FROM_EXTENSION:
399 extension = pathlib.Path(local_path).suffix.lower()
400 if extension in so_compression.get_supported_extensions():
401 return None
402 elif compression != so_compression.NO_COMPRESSION:
403 return None
405 open_kwargs: dict[str, Any] = {}
406 if encoding is not None:
407 open_kwargs["encoding"] = encoding
408 mode = mode.replace("b", "")
409 if newline is not None:
410 open_kwargs["newline"] = newline
412 #
413 # binary mode of the builtin/stdlib open function doesn't take an errors argument
414 #
415 if errors and "b" not in mode:
416 open_kwargs["errors"] = errors
418 return _builtin_open(local_path, mode, buffering=buffering, **open_kwargs)
421def _open_binary_stream(uri: Uri, mode: str, transport_params: TransportParams) -> IO[bytes]:
422 """Open an arbitrary URI in the specified binary mode.
424 Not all modes are supported for all protocols.
426 Args:
427 uri: The URI to open. May be a string, or something else.
428 mode: The mode to open with. Must be rb, wb or ab.
429 transport_params: Keyword arguments for the transport layer.
431 Returns:
432 A file-like object with a ``.name`` attribute.
434 Raises:
435 NotImplementedError: If ``mode`` is not a supported binary mode.
436 TypeError: If ``uri`` is not a string or integer file descriptor.
437 """
438 if mode not in ("rb", "rb+", "wb", "wb+", "ab", "ab+"):
439 #
440 # This should really be a ValueError, but for the sake of compatibility
441 # with older versions, which raise NotImplementedError, we do the same.
442 #
443 msg = f"unsupported mode: {mode!r}"
444 raise NotImplementedError(msg)
446 if isinstance(uri, int):
447 #
448 # We're working with a file descriptor. If we open it, its name is
449 # just the integer value, which isn't helpful. Unfortunately, there's
450 # no easy cross-platform way to go from a file descriptor to the filename,
451 # so we just give up here. The user will have to handle their own
452 # compression, etc. explicitly.
453 #
454 return _builtin_open(uri, mode, closefd=False)
456 if not isinstance(uri, str):
457 msg = f"don't know how to handle uri {uri!r}"
458 raise TypeError(msg)
460 scheme = _sniff_scheme(uri)
461 submodule = transport.get_transport(scheme)
462 fobj = submodule.open_uri(uri, mode, transport_params)
463 if not hasattr(fobj, "name"):
464 fobj.name = uri
466 return fobj
469def _encoding_wrapper(
470 fileobj: IO[Any],
471 mode: str,
472 encoding: str | None = None,
473 errors: str | None = None,
474 newline: str | None = None,
475) -> IO[Any]:
476 """Decode bytes into text, if necessary.
478 If mode specifies binary access, does nothing, unless the encoding is
479 specified. A non-null encoding implies text mode.
481 Args:
482 fileobj: Must quack like a filehandle object.
483 mode: The mode which was originally requested by the user.
484 encoding: The text encoding to use. If mode is binary, overrides mode.
485 errors: The method to use when handling encoding/decoding errors.
486 newline: Forwarded to the text wrapper.
488 Returns:
489 A file object.
490 """
491 #
492 # If the mode is binary, but the user specified an encoding, assume they
493 # want text. If we don't make this assumption, ignore the encoding and
494 # return bytes, smart_open behavior will diverge from the built-in open:
495 #
496 # open(filename, encoding='utf-8') returns a text stream in Py3
497 # smart_open(filename, encoding='utf-8') would return a byte stream
498 # without our assumption, because the default mode is rb.
499 #
500 if "b" in mode and encoding is None:
501 return fileobj
503 if encoding is None:
504 encoding = DEFAULT_ENCODING
506 return so_utils.TextIOWrapper(
507 fileobj,
508 encoding=encoding,
509 errors=errors,
510 newline=newline,
511 write_through=True,
512 )
515class patch_pathlib: # noqa: N801 # function-shaped name in public API
516 """Replace `Path.open` with `smart_open.open`."""
518 def __init__(self) -> None:
519 self.old_impl = _patch_pathlib(open)
521 def __enter__(self) -> Self: # noqa: D105
522 return self
524 def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: # noqa: D105
525 _patch_pathlib(self.old_impl)
528def _patch_pathlib(func: Callable[..., Any]) -> Callable[..., Any]:
529 """Replace `Path.open` with `func`."""
530 old_impl = pathlib.Path.open
531 pathlib.Path.open = func # ty: ignore[invalid-assignment] # intentional monkeypatch
532 return old_impl
535#
536# Prevent failures with doctools from messing up the entire library. We don't
537# expect such failures, but contributed modules (e.g. new transport mechanisms)
538# may not be as polished.
539#
540try:
541 doctools.tweak_open_docstring(open)
542 doctools.tweak_parse_uri_docstring(parse_uri)
543except Exception:
544 logger.exception(
545 "Encountered a non-fatal error while building docstrings (see below). "
546 "help(smart_open) will provide incomplete information as a result. "
547 "For full help text, see "
548 "<https://github.com/piskvorky/smart_open/blob/master/help.txt>."
549 )