Skip to content

Commit 43c9b2c

Browse files
kartikgupta-dbkartikgupta-db
authored
Make dbutils typecast use a valid type variable (#259)
## Changes * We used an variable for collecting type information about dbutils. This behaviour was patched in a recent release of pyright (microsoft/pyright@6169e0f). Now we need to use an explicit type variable. ## Tests - [x] tested manually on vscode with pylance `v2023.7.40` and type checking mode set to `strict` - [x] `make test` run locally - [x] `make fmt` applied - [x] relevant integration tests applied
1 parent 519fbe3 commit 43c9b2c

File tree

3 files changed

+369
-367
lines changed

3 files changed

+369
-367
lines changed

‎databricks/sdk/runtime/__init__.py‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
from __future__ import annotations
22

3+
from typing import Union
4+
35
is_local_implementation = True
46

57
# All objects that are injected into the Notebook's user namespace should also be made
@@ -33,21 +35,22 @@
3335
_globals[var] = userNamespaceGlobals[var]
3436
is_local_implementation = False
3537
except ImportError:
36-
from typing import Type, cast
38+
from typing import cast
3739

3840
# OSS implementation
3941
is_local_implementation = True
4042

43+
from databricks.sdk.dbutils import RemoteDbUtils
44+
45+
from . import dbutils_stub
46+
47+
dbutils_type = Union[dbutils_stub.dbutils, RemoteDbUtils]
48+
4149
try:
42-
from . import stub
4350
from .stub import *
44-
dbutils_type = Type[stub.dbutils]
4551
except (ImportError, NameError):
46-
from databricks.sdk.dbutils import RemoteDbUtils
47-
4852
# this assumes that all environment variables are set
4953
dbutils = RemoteDbUtils()
50-
dbutils_type = RemoteDbUtils
5154

5255
dbutils = cast(dbutils_type, dbutils)
5356

‎databricks/sdk/runtime/dbutils_stub.py‎

Lines changed: 360 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,360 @@
1+
import typing
2+
from collections import namedtuple
3+
4+
5+
class FileInfo(namedtuple('FileInfo', ['path', 'name', 'size', "modificationTime"])):
6+
pass
7+
8+
9+
class MountInfo(namedtuple('MountInfo', ['mountPoint', 'source', 'encryptionType'])):
10+
pass
11+
12+
13+
class SecretScope(namedtuple('SecretScope', ['name'])):
14+
15+
def getName(self):
16+
return self.name
17+
18+
19+
class SecretMetadata(namedtuple('SecretMetadata', ['key'])):
20+
pass
21+
22+
23+
class dbutils:
24+
25+
class credentials:
26+
"""
27+
Utilities for interacting with credentials within notebooks
28+
"""
29+
30+
@staticmethod
31+
def assumeRole(role: str) -> bool:
32+
"""
33+
Sets the role ARN to assume when looking for credentials to authenticate with S3
34+
"""
35+
...
36+
37+
@staticmethod
38+
def showCurrentRole() -> typing.List[str]:
39+
"""
40+
Shows the currently set role
41+
"""
42+
...
43+
44+
@staticmethod
45+
def showRoles() -> typing.List[str]:
46+
"""
47+
Shows the set of possibly assumed roles
48+
"""
49+
...
50+
51+
@staticmethod
52+
def getCurrentCredentials() -> typing.Mapping[str, str]:
53+
...
54+
55+
class data:
56+
"""
57+
Utilities for understanding and interacting with datasets (EXPERIMENTAL)
58+
"""
59+
60+
@staticmethod
61+
def summarize(df: any, precise: bool = False) -> None:
62+
"""Summarize a Spark/pandas/Koalas DataFrame and visualize the statistics to get quick insights.
63+
64+
Example: dbutils.data.summarize(df)
65+
66+
:param df: A pyspark.sql.DataFrame, pyspark.pandas.DataFrame, databricks.koalas.DataFrame
67+
or pandas.DataFrame object to summarize. Streaming dataframes are not supported.
68+
:param precise: If false, percentiles, distinct item counts, and frequent item counts
69+
will be computed approximately to reduce the run time.
70+
If true, distinct item counts and frequent item counts will be computed exactly,
71+
and percentiles will be computed with high precision.
72+
73+
:return: visualization of the computed summmary statistics.
74+
"""
75+
...
76+
77+
class fs:
78+
"""
79+
Manipulates the Databricks filesystem (DBFS) from the console
80+
"""
81+
82+
@staticmethod
83+
def cp(source: str, dest: str, recurse: bool = False) -> bool:
84+
"""
85+
Copies a file or directory, possibly across FileSystems
86+
"""
87+
...
88+
89+
@staticmethod
90+
def head(file: str, max_bytes: int = 65536) -> str:
91+
"""
92+
Returns up to the first 'maxBytes' bytes of the given file as a String encoded in UTF-8
93+
"""
94+
...
95+
96+
@staticmethod
97+
def ls(path: str) -> typing.List[FileInfo]:
98+
"""
99+
Lists the contents of a directory
100+
"""
101+
...
102+
103+
@staticmethod
104+
def mkdirs(dir: str) -> bool:
105+
"""
106+
Creates the given directory if it does not exist, also creating any necessary parent directories
107+
"""
108+
...
109+
110+
@staticmethod
111+
def mv(source: str, dest: str, recurse: bool = False) -> bool:
112+
"""
113+
Moves a file or directory, possibly across FileSystems
114+
"""
115+
...
116+
117+
@staticmethod
118+
def put(file: str, contents: str, overwrite: bool = False) -> bool:
119+
"""
120+
Writes the given String out to a file, encoded in UTF-8
121+
"""
122+
...
123+
124+
@staticmethod
125+
def rm(dir: str, recurse: bool = False) -> bool:
126+
"""
127+
Removes a file or directory
128+
"""
129+
...
130+
131+
@staticmethod
132+
def cacheFiles(*files):
133+
...
134+
135+
@staticmethod
136+
def cacheTable(name: str):
137+
...
138+
139+
@staticmethod
140+
def uncacheFiles(*files):
141+
...
142+
143+
@staticmethod
144+
def uncacheTable(name: str):
145+
...
146+
147+
@staticmethod
148+
def mount(source: str,
149+
mount_point: str,
150+
encryption_type: str = "",
151+
owner: typing.Optional[str] = None,
152+
extra_configs: typing.Mapping[str, str] = {},
153+
) -> bool:
154+
"""
155+
Mounts the given source directory into DBFS at the given mount point
156+
"""
157+
...
158+
159+
@staticmethod
160+
def updateMount(source: str,
161+
mount_point: str,
162+
encryption_type: str = "",
163+
owner: typing.Optional[str] = None,
164+
extra_configs: typing.Mapping[str, str] = {},
165+
) -> bool:
166+
"""
167+
Similar to mount(), but updates an existing mount point (if present) instead of creating a new one
168+
"""
169+
...
170+
171+
@staticmethod
172+
def mounts() -> typing.List[MountInfo]:
173+
"""
174+
Displays information about what is mounted within DBFS
175+
"""
176+
...
177+
178+
@staticmethod
179+
def refreshMounts() -> bool:
180+
"""
181+
Forces all machines in this cluster to refresh their mount cache, ensuring they receive the most recent information
182+
"""
183+
...
184+
185+
@staticmethod
186+
def unmount(mount_point: str) -> bool:
187+
"""
188+
Deletes a DBFS mount point
189+
"""
190+
...
191+
192+
class jobs:
193+
"""
194+
Utilities for leveraging jobs features
195+
"""
196+
197+
class taskValues:
198+
"""
199+
Provides utilities for leveraging job task values
200+
"""
201+
202+
@staticmethod
203+
def get(taskKey: str, key: str, default: any = None, debugValue: any = None) -> None:
204+
"""
205+
Returns the latest task value that belongs to the current job run
206+
"""
207+
...
208+
209+
@staticmethod
210+
def set(key: str, value: any) -> None:
211+
"""
212+
Sets a task value on the current task run
213+
"""
214+
...
215+
216+
class library:
217+
"""
218+
Utilities for session isolated libraries
219+
"""
220+
221+
@staticmethod
222+
def restartPython() -> None:
223+
"""
224+
Restart python process for the current notebook session
225+
"""
226+
...
227+
228+
class notebook:
229+
"""
230+
Utilities for the control flow of a notebook (EXPERIMENTAL)
231+
"""
232+
233+
@staticmethod
234+
def exit(value: str) -> None:
235+
"""
236+
This method lets you exit a notebook with a value
237+
"""
238+
...
239+
240+
@staticmethod
241+
def run(path: str, timeout_seconds: int, arguments: typing.Mapping[str, str]) -> str:
242+
"""
243+
This method runs a notebook and returns its exit value
244+
"""
245+
...
246+
247+
class secrets:
248+
"""
249+
Provides utilities for leveraging secrets within notebooks
250+
"""
251+
252+
@staticmethod
253+
def get(scope: str, key: str) -> str:
254+
"""
255+
Gets the string representation of a secret value with scope and key
256+
"""
257+
...
258+
259+
@staticmethod
260+
def getBytes(self, scope: str, key: str) -> bytes:
261+
"""Gets the bytes representation of a secret value for the specified scope and key."""
262+
263+
@staticmethod
264+
def list(scope: str) -> typing.List[SecretMetadata]:
265+
"""
266+
Lists secret metadata for secrets within a scope
267+
"""
268+
...
269+
270+
@staticmethod
271+
def listScopes() -> typing.List[SecretScope]:
272+
"""
273+
Lists secret scopes
274+
"""
275+
...
276+
277+
class widgets:
278+
"""
279+
provides utilities for working with notebook widgets. You can create different types of widgets and get their bound value
280+
"""
281+
282+
@staticmethod
283+
def get(name: str) -> str:
284+
"""Returns the current value of a widget with give name.
285+
:param name: Name of the argument to be accessed
286+
:return: Current value of the widget or default value
287+
"""
288+
...
289+
290+
@staticmethod
291+
def getArgument(name: str, defaultValue: typing.Optional[str] = None) -> str:
292+
"""Returns the current value of a widget with give name.
293+
:param name: Name of the argument to be accessed
294+
:param defaultValue: (Deprecated) default value
295+
:return: Current value of the widget or default value
296+
"""
297+
...
298+
299+
@staticmethod
300+
def text(name: str, defaultValue: str, label: str = None):
301+
"""Creates a text input widget with given name, default value and optional label for
302+
display
303+
:param name: Name of argument associated with the new input widget
304+
:param defaultValue: Default value of the input widget
305+
:param label: Optional label string for display in notebook and dashboard
306+
"""
307+
...
308+
309+
@staticmethod
310+
def dropdown(name: str, defaultValue: str, choices: typing.List[str], label: str = None):
311+
"""Creates a dropdown input widget with given specification.
312+
:param name: Name of argument associated with the new input widget
313+
:param defaultValue: Default value of the input widget (must be one of choices)
314+
:param choices: List of choices for the dropdown input widget
315+
:param label: Optional label string for display in notebook and dashboard
316+
"""
317+
...
318+
319+
@staticmethod
320+
def combobox(name: str,
321+
defaultValue: str,
322+
choices: typing.List[str],
323+
label: typing.Optional[str] = None,
324+
):
325+
"""Creates a combobox input widget with given specification.
326+
:param name: Name of argument associated with the new input widget
327+
:param defaultValue: Default value of the input widget
328+
:param choices: List of choices for the dropdown input widget
329+
:param label: Optional label string for display in notebook and dashboard
330+
"""
331+
...
332+
333+
@staticmethod
334+
def multiselect(name: str,
335+
defaultValue: str,
336+
choices: typing.List[str],
337+
label: typing.Optional[str] = None,
338+
):
339+
"""Creates a multiselect input widget with given specification.
340+
:param name: Name of argument associated with the new input widget
341+
:param defaultValue: Default value of the input widget (must be one of choices)
342+
:param choices: List of choices for the dropdown input widget
343+
:param label: Optional label string for display in notebook and dashboard
344+
"""
345+
...
346+
347+
@staticmethod
348+
def remove(name: str):
349+
"""Removes given input widget. If widget does not exist it will throw an error.
350+
:param name: Name of argument associated with input widget to be removed
351+
"""
352+
...
353+
354+
@staticmethod
355+
def removeAll():
356+
"""Removes all input widgets in the notebook."""
357+
...
358+
359+
360+
getArgument = dbutils.widgets.getArgument

0 commit comments

Comments
 (0)