Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/smart_open/transport.py: 84%

1# 

2# Copyright (C) 2020 Radim Rehurek <me@radimrehurek.com> 

3# 

4# This code is distributed under the terms and conditions 

5# from the MIT License (MIT). 

6# 

7"""Maintains a registry of transport mechanisms. 

8 

9The main entrypoint is :func:`get_transport`. See also :file:`EXTENDING.md`. 

10 

11""" 

12 

13from __future__ import annotations 

14 

15import importlib 

16import logging 

17from typing import TYPE_CHECKING 

18 

19import smart_open.local_file 

20 

21if TYPE_CHECKING: 

22 from types import ModuleType 

23 

24logger = logging.getLogger(__name__) 

25 

26NO_SCHEME = "" 

27 

28_REGISTRY: dict[str, ModuleType] = {NO_SCHEME: smart_open.local_file} 

29_ERRORS: dict[str, str] = {} 

30_MISSING_DEPS_ERROR = """You are trying to use the %(module)s functionality of smart_open 

31but you do not have the correct %(module)s dependencies installed. Try: 

32 

33 pip install smart_open[%(module)s] 

34 

35""" 

36 

37 

38def register_transport(submodule: str | ModuleType) -> None: 

39 """Register a submodule as a transport mechanism for ``smart_open``. 

40 

41 This module **must** have: 

42 

43 - `SCHEME` attribute (or `SCHEMES`, if the submodule supports multiple schemes) 

44 - `open` function 

45 - `open_uri` function 

46 - `parse_uri' function 

47 

48 Once registered, you can get the submodule by calling :func:`get_transport`. 

49 

50 """ 

51 if isinstance(submodule, str): 

52 module_name = submodule 

53 try: 

54 submodule = importlib.import_module(submodule) 

55 except ImportError: 

56 return 

57 else: 

58 module_name = submodule.__name__ 

59 # Save only the last module name piece 

60 module_name = module_name.rsplit(".")[-1] 

61 

62 if hasattr(submodule, "SCHEME"): 

63 schemes = [submodule.SCHEME] 

64 elif hasattr(submodule, "SCHEMES"): 

65 schemes = submodule.SCHEMES 

66 else: 

67 msg = f"{submodule!r} does not have a .SCHEME or .SCHEMES attribute" 

68 raise ValueError(msg) 

69 

70 for f in ("open", "open_uri", "parse_uri"): 

71 assert hasattr(submodule, f), f"{submodule!r} is missing {f!r}" # noqa: S101 # internal precondition; misuse should crash loudly 

72 

73 for scheme in schemes: 

74 assert scheme not in _REGISTRY # noqa: S101 # internal precondition; misuse should crash loudly 

75 if getattr(submodule, "MISSING_DEPS", False): 

76 _ERRORS[scheme] = module_name 

77 else: 

78 _REGISTRY[scheme] = submodule 

79 

80 

81def get_transport(scheme: str) -> ModuleType: 

82 """Get the submodule that handles transport for the specified scheme. 

83 

84 This submodule must have been previously registered via :func:`register_transport`. 

85 

86 """ 

87 expected = SUPPORTED_SCHEMES 

88 readme_url = "https://github.com/piskvorky/smart_open/blob/master/README.md" 

89 message = ( 

90 "Unable to handle scheme {scheme!r}, expected one of {expected!r}. " 

91 "Extra dependencies required by {scheme!r} may be missing. " 

92 "See <{readme_url}> for details.".format(**locals()) 

93 ) 

94 if scheme in _ERRORS: 

95 raise ImportError(_MISSING_DEPS_ERROR % {"module": _ERRORS[scheme]}) 

96 if scheme in _REGISTRY: 

97 return _REGISTRY[scheme] 

98 raise NotImplementedError(message) 

99 

100 

101register_transport(smart_open.local_file) 

102register_transport("smart_open.azure") 

103register_transport("smart_open.ftp") 

104register_transport("smart_open.gcs") 

105register_transport("smart_open.hdfs") 

106register_transport("smart_open.http") 

107register_transport("smart_open.s3") 

108register_transport("smart_open.ssh") 

109register_transport("smart_open.webhdfs") 

110 

111SUPPORTED_SCHEMES = tuple(sorted(_REGISTRY.keys())) 

112"""The transport schemes that the local installation of ``smart_open`` supports."""