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", "x", "r+", "w+", "a+", "rt", "wt", "at", "xt"] = ...,
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", "xb", "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: FBT001, FBT002 # public API
144 opener: Callable[[str, int], int] | None = None,
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 logger.debug("%r", locals())
208 if not isinstance(mode, str):
209 msg = "mode should be a string"
210 raise TypeError(msg)
212 if compression not in so_compression.get_supported_compression_types():
213 msg = f"invalid compression type: {compression}"
214 raise ValueError(msg)
216 if transport_params is None:
217 transport_params = {}
219 fobj = _shortcut_open(
220 uri,
221 mode,
222 compression=compression,
223 buffering=buffering,
224 encoding=encoding,
225 errors=errors,
226 newline=newline,
227 )
228 if fobj is not None:
229 return fobj
231 #
232 # This is a work-around for the problem described in Issue #144.
233 # If the user has explicitly specified an encoding, then assume they want
234 # us to open the destination in text mode, instead of the default binary.
235 #
236 # If we change the default mode to be text, and match the normal behavior
237 # of Py2 and 3, then the above assumption will be unnecessary.
238 #
239 if encoding is not None and "b" in mode:
240 mode = mode.replace("b", "")
242 if isinstance(uri, pathlib.Path):
243 uri = str(uri)
245 explicit_encoding = encoding
246 encoding = explicit_encoding or DEFAULT_ENCODING
248 #
249 # This is how we get from the filename to the end result. Decompression is
250 # optional, but it always accepts bytes and returns bytes.
251 #
252 # Decoding is also optional, accepts bytes and returns text. The diagram
253 # below is for reading, for writing, the flow is from right to left, but
254 # the code is identical.
255 #
256 # open as binary decompress? decode?
257 # filename ---------------> bytes -------------> bytes ---------> text
258 # binary decompressed decode
259 #
261 try:
262 binary_mode = _get_binary_mode(mode)
263 except ValueError as ve:
264 raise NotImplementedError(ve.args[0]) from ve
266 binary = _open_binary_stream(uri, binary_mode, transport_params)
267 name = getattr(binary, "name", None)
268 # prefer the stream's own name; if it's not string-like (e.g. ftp socket fileno), fall back to uri
269 filename = name if isinstance(name, str) else uri if isinstance(uri, str) else None
270 decompressed = so_compression.compression_wrapper(
271 binary,
272 binary_mode,
273 compression,
274 filename=filename,
275 compression_kwargs=compression_kwargs,
276 )
278 if "b" not in mode or explicit_encoding is not None:
279 decoded = _encoding_wrapper(
280 decompressed,
281 mode,
282 encoding=encoding,
283 errors=errors,
284 newline=newline,
285 )
286 else:
287 decoded = decompressed
289 #
290 # There are some useful methods in the binary readers, e.g. to_boto3, that get
291 # hidden by the multiple layers of wrapping we just performed. Promote
292 # them so they are visible to the user.
293 #
294 if decoded != binary:
295 promoted_attrs = ["to_boto3"]
296 for attr in promoted_attrs:
297 with contextlib.suppress(AttributeError):
298 setattr(decoded, attr, getattr(binary, attr))
300 return cast("IO[Any]", so_utils.FileLikeProxy(decoded, binary))
303def _get_binary_mode(mode_str: str) -> str: # noqa: C901 # legacy internal helper; refactor in a dedicated PR
304 #
305 # https://docs.python.org/3/library/functions.html#open
306 #
307 # The order of characters in the mode parameter appears to be unspecified.
308 # The implementation follows the examples, just to be safe.
309 #
310 mode = list(mode_str)
311 binmode = []
313 if "t" in mode and "b" in mode:
314 msg = "can't have text and binary mode at once"
315 raise ValueError(msg)
317 counts = [mode.count(x) for x in "rwa"]
318 if sum(counts) > 1:
319 msg = "must have exactly one of create/read/write/append mode"
320 raise ValueError(msg)
322 def transfer(char: str) -> None:
323 binmode.append(mode.pop(mode.index(char)))
325 if "a" in mode:
326 transfer("a")
327 elif "w" in mode:
328 transfer("w")
329 elif "r" in mode:
330 transfer("r")
331 else:
332 msg = "Must have exactly one of create/read/write/append mode and at most one plus"
333 raise ValueError(msg)
335 if "b" in mode:
336 transfer("b")
337 elif "t" in mode:
338 mode.pop(mode.index("t"))
339 binmode.append("b")
340 else:
341 binmode.append("b")
343 if "+" in mode:
344 transfer("+")
346 #
347 # There shouldn't be anything left in the mode list at this stage.
348 # If there is, then either we've missed something and the implementation
349 # of this function is broken, or the original input mode is invalid.
350 #
351 if mode:
352 msg = f"invalid mode: {mode_str!r}"
353 raise ValueError(msg)
355 return "".join(binmode)
358def _shortcut_open( # noqa: PLR0913 # legacy internal helper; refactor in a dedicated PR
359 uri: Uri,
360 mode: str,
361 compression: str,
362 buffering: int = -1,
363 encoding: str | None = None,
364 errors: str | None = None,
365 newline: str | None = None,
366) -> IO[Any] | None:
367 """Try to open the URI using the standard library io.open function.
369 This can be much faster than the alternative of opening in binary mode and
370 then decoding.
372 This is only possible under the following conditions:
374 1. Opening a local file; and
375 2. Compression is disabled
377 If it is not possible to use the built-in open for the specified URI,
378 returns None.
380 Args:
381 uri: A string indicating what to open.
382 mode: The mode to pass to the open function.
383 compression: The compression type selected.
384 buffering: Mimics built-in open parameter of the same name.
385 encoding: Mimics built-in open parameter of the same name.
386 errors: Mimics built-in open parameter of the same name.
387 newline: Mimics built-in open parameter of the same name.
389 Returns:
390 The opened file, or None if no shortcut is possible.
391 """
392 if not isinstance(uri, str):
393 return None
395 scheme = _sniff_scheme(uri)
396 if scheme not in (transport.NO_SCHEME, so_file.SCHEME):
397 return None
399 local_path = so_file.extract_local_path(uri)
400 if compression == so_compression.INFER_FROM_EXTENSION:
401 extension = pathlib.Path(local_path).suffix
402 if extension in so_compression.get_supported_extensions():
403 return None
404 elif compression != so_compression.NO_COMPRESSION:
405 return None
407 open_kwargs: dict[str, Any] = {}
408 if encoding is not None:
409 open_kwargs["encoding"] = encoding
410 mode = mode.replace("b", "")
411 if newline is not None:
412 open_kwargs["newline"] = newline
414 #
415 # binary mode of the builtin/stdlib open function doesn't take an errors argument
416 #
417 if errors and "b" not in mode:
418 open_kwargs["errors"] = errors
420 return _builtin_open(local_path, mode, buffering=buffering, **open_kwargs)
423def _open_binary_stream(uri: Uri, mode: str, transport_params: TransportParams) -> IO[bytes]:
424 """Open an arbitrary URI in the specified binary mode.
426 Not all modes are supported for all protocols.
428 Args:
429 uri: The URI to open. May be a string, or something else.
430 mode: The mode to open with. Must be rb, wb or ab.
431 transport_params: Keyword arguments for the transport layer.
433 Returns:
434 A file-like object with a ``.name`` attribute.
436 Raises:
437 NotImplementedError: If ``mode`` is not a supported binary mode.
438 TypeError: If ``uri`` is not a string or integer file descriptor.
439 """
440 if mode not in ("rb", "rb+", "wb", "wb+", "ab", "ab+"):
441 #
442 # This should really be a ValueError, but for the sake of compatibility
443 # with older versions, which raise NotImplementedError, we do the same.
444 #
445 msg = f"unsupported mode: {mode!r}"
446 raise NotImplementedError(msg)
448 if isinstance(uri, int):
449 #
450 # We're working with a file descriptor. If we open it, its name is
451 # just the integer value, which isn't helpful. Unfortunately, there's
452 # no easy cross-platform way to go from a file descriptor to the filename,
453 # so we just give up here. The user will have to handle their own
454 # compression, etc. explicitly.
455 #
456 return _builtin_open(uri, mode, closefd=False)
458 if not isinstance(uri, str):
459 msg = f"don't know how to handle uri {uri!r}"
460 raise TypeError(msg)
462 scheme = _sniff_scheme(uri)
463 submodule = transport.get_transport(scheme)
464 fobj = submodule.open_uri(uri, mode, transport_params)
465 if not hasattr(fobj, "name"):
466 fobj.name = uri
468 return fobj
471def _encoding_wrapper(
472 fileobj: IO[Any],
473 mode: str,
474 encoding: str | None = None,
475 errors: str | None = None,
476 newline: str | None = None,
477) -> IO[Any]:
478 """Decode bytes into text, if necessary.
480 If mode specifies binary access, does nothing, unless the encoding is
481 specified. A non-null encoding implies text mode.
483 Args:
484 fileobj: Must quack like a filehandle object.
485 mode: The mode which was originally requested by the user.
486 encoding: The text encoding to use. If mode is binary, overrides mode.
487 errors: The method to use when handling encoding/decoding errors.
488 newline: Forwarded to the text wrapper.
490 Returns:
491 A file object.
492 """
493 logger.debug("encoding_wrapper: %r", locals())
495 #
496 # If the mode is binary, but the user specified an encoding, assume they
497 # want text. If we don't make this assumption, ignore the encoding and
498 # return bytes, smart_open behavior will diverge from the built-in open:
499 #
500 # open(filename, encoding='utf-8') returns a text stream in Py3
501 # smart_open(filename, encoding='utf-8') would return a byte stream
502 # without our assumption, because the default mode is rb.
503 #
504 if "b" in mode and encoding is None:
505 return fileobj
507 if encoding is None:
508 encoding = DEFAULT_ENCODING
510 return so_utils.TextIOWrapper(
511 fileobj,
512 encoding=encoding,
513 errors=errors,
514 newline=newline,
515 write_through=True,
516 )
519class patch_pathlib: # noqa: N801 # function-shaped name in public API
520 """Replace `Path.open` with `smart_open.open`."""
522 def __init__(self) -> None:
523 self.old_impl = _patch_pathlib(open)
525 def __enter__(self) -> Self: # noqa: D105
526 return self
528 def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> None: # noqa: D105
529 _patch_pathlib(self.old_impl)
532def _patch_pathlib(func: Callable[..., Any]) -> Callable[..., Any]:
533 """Replace `Path.open` with `func`."""
534 old_impl = pathlib.Path.open
535 pathlib.Path.open = func # ty: ignore[invalid-assignment] # intentional monkeypatch
536 return old_impl
539#
540# Prevent failures with doctools from messing up the entire library. We don't
541# expect such failures, but contributed modules (e.g. new transport mechanisms)
542# may not be as polished.
543#
544try:
545 doctools.tweak_open_docstring(open)
546 doctools.tweak_parse_uri_docstring(parse_uri)
547except Exception:
548 logger.exception(
549 "Encountered a non-fatal error while building docstrings (see below). "
550 "help(smart_open) will provide incomplete information as a result. "
551 "For full help text, see "
552 "<https://github.com/piskvorky/smart_open/blob/master/help.txt>."
553 )