Repository navigation
Expand file tree
/
Copy path__init__.py
More file actions
142 lines (112 loc) · 4.36 KB
/
Copy path__init__.py
File metadata and controls
142 lines (112 loc) · 4.36 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
"""Initialize Ibis module."""
from __future__ import annotations
__version__ = "10.6.0"
import warnings
from typing import Any
from ibis import examples, util
from ibis.backends import BaseBackend
from ibis.common.exceptions import IbisError
from ibis.config import options
from ibis.expr import api
from ibis.expr import types as ir
from ibis.expr.api import * # noqa: F403
from ibis.expr.operations import udf
__all__ = [ # noqa: PLE0604
"api",
"examples",
"ir",
"udf",
"util",
"BaseBackend",
"IbisError",
"options",
*api.__all__,
]
_KNOWN_BACKENDS = ["heavyai"]
def __dir__() -> list[str]:
"""Adds tab completion for ibis backends to the top-level module."""
out = set(__all__)
out.update(ep.name for ep in util.backend_entry_points())
return sorted(out)
def load_backend(name: str) -> BaseBackend:
"""Load backends in a lazy way with `ibis.<backend-name>`.
This also registers the backend options.
Examples
--------
>>> import ibis
>>> con = ibis.sqlite.connect(...)
When accessing the `sqlite` attribute of the `ibis` module, this function
is called, and a backend with the `sqlite` name is tried to load from
the `ibis.backends` entrypoints. If successful, the `ibis.sqlite`
attribute is "cached", so this function is only called the first time.
"""
entry_points = {ep for ep in util.backend_entry_points() if ep.name == name}
if not entry_points:
msg = f"module 'ibis' has no attribute '{name}'. "
if name in _KNOWN_BACKENDS:
msg += f"""If you are trying to access the '{name}' backend,
try installing it first with `pip install 'ibis-framework[{name}]'`"""
raise AttributeError(msg)
if len(entry_points) > 1:
raise RuntimeError(
f"{len(entry_points)} packages found for backend '{name}': "
f"{entry_points}\n"
"There should be only one, please uninstall the unused packages "
"and just leave the one that needs to be used."
)
import types
import ibis
(entry_point,) = entry_points
try:
module = entry_point.load()
except ImportError as exc:
raise ImportError(
f"Failed to import the {name} backend due to missing dependencies.\n\n"
f"You can pip or conda install the {name} backend as follows:\n\n"
f' python -m pip install -U "ibis-framework[{name}]" # pip install\n'
f" conda install -c conda-forge ibis-{name} # or conda install"
) from exc
backend = module.Backend()
# The first time a backend is loaded, we register its options, and we set
# it as an attribute of `ibis`, so `__getattr__` is not called again for it
backend.register_options()
# We don't want to expose all the methods on an unconnected backend to the user.
# In lieu of a full redesign, we create a proxy module and add only the methods
# that are valid to call without a connect call. These are:
#
# - connect
# - compile
# - has_operation
# - _from_url
#
# We also copy over the docstring from `do_connect` to the proxy `connect`
# method, since that's where all the backend-specific kwargs are currently
# documented. This is all admittedly gross, but it works and doesn't
# require a backend redesign yet.
def connect(*args, **kwargs):
return backend.connect(*args, **kwargs)
connect.__doc__ = backend.do_connect.__doc__
connect.__wrapped__ = backend.do_connect
connect.__module__ = f"ibis.{name}"
proxy = types.ModuleType(f"ibis.{name}")
setattr(ibis, name, proxy)
proxy.connect = connect
proxy.compile = backend.compile
proxy.has_operation = backend.has_operation
proxy.name = name
proxy._from_url = backend._from_url
# Add any additional methods that should be exposed at the top level
for attr in getattr(backend, "_top_level_methods", ()):
setattr(proxy, attr, getattr(backend, attr))
return proxy
def __getattr__(name: str) -> Any:
if name == "NA":
warnings.warn(
"The 'ibis.NA' constant is deprecated as of v9.1 and will be removed in a future "
"version. Use 'ibis.null()' instead.",
DeprecationWarning,
stacklevel=2,
)
return null() # noqa: F405
else:
return load_backend(name)