Skip to content

Commit e627fee

Browse files
authored
Fixed KeyError from assess_workflows task (#2919)
## Changes Log when inferring node results in `KeyError`. ### Linked issues Fixes #2908 ### Functionality - [x] modified existing workflow: `assessment`
1 parent 38a5df2 commit e627fee

File tree

2 files changed

+17
-17
lines changed

2 files changed

+17
-17
lines changed

src/databricks/labs/ucx/source_code/python/python_infer.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -66,15 +66,15 @@ def _infer_values(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
6666
yield from cls._safe_infer_internal(node)
6767

6868
@classmethod
69-
def _safe_infer_internal(cls, node: NodeNG):
69+
def _safe_infer_internal(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
7070
try:
7171
yield from cls._unsafe_infer_internal(node)
72-
except InferenceError as e:
73-
logger.debug(f"When inferring {node}", exc_info=e)
72+
except (InferenceError, KeyError) as e:
73+
logger.debug(f"When inferring: {node}", exc_info=e)
7474
yield [Uninferable]
7575

7676
@classmethod
77-
def _unsafe_infer_internal(cls, node: NodeNG):
77+
def _unsafe_infer_internal(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
7878
all_inferred = node.inferred()
7979
if len(all_inferred) == 0 and isinstance(node, Instance):
8080
yield [node]
@@ -148,7 +148,7 @@ def infer_call_result(self, context: InferenceContext | None = None, **_): # ca
148148
class _LocalInferredValue(InferredValue):
149149

150150
@classmethod
151-
def do_infer_values(cls, node: NodeNG):
151+
def do_infer_values(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
152152
yield from cls._infer_values(node)
153153

154154

tests/unit/source_code/python/test_python_infer.py

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,31 +5,31 @@
55
from databricks.labs.ucx.source_code.python.python_infer import InferredValue
66

77

8-
def test_infers_empty_list():
8+
def test_infers_empty_list() -> None:
99
tree = Tree.parse("a=[]")
1010
nodes = tree.locate(Assign, [])
1111
tree = Tree(nodes[0].value)
1212
values = list(InferredValue.infer_from_node(tree.node))
1313
assert not values
1414

1515

16-
def test_infers_empty_tuple():
16+
def test_infers_empty_tuple() -> None:
1717
tree = Tree.parse("a=tuple()")
1818
nodes = tree.locate(Assign, [])
1919
tree = Tree(nodes[0].value)
2020
values = list(InferredValue.infer_from_node(tree.node))
2121
assert not values
2222

2323

24-
def test_infers_empty_set():
24+
def test_infers_empty_set() -> None:
2525
tree = Tree.parse("a={}")
2626
nodes = tree.locate(Assign, [])
2727
tree = Tree(nodes[0].value)
2828
values = list(InferredValue.infer_from_node(tree.node))
2929
assert not values
3030

3131

32-
def test_infers_fstring_value():
32+
def test_infers_fstring_value() -> None:
3333
source = """
3434
value = "abc"
3535
fstring = f"Hello {value}!"
@@ -43,7 +43,7 @@ def test_infers_fstring_value():
4343
assert strings == ["Hello abc!"]
4444

4545

46-
def test_infers_fstring_dict_value():
46+
def test_infers_fstring_dict_value() -> None:
4747
source = """
4848
value = { "abc": 123 }
4949
fstring = f"Hello {value['abc']}!"
@@ -57,7 +57,7 @@ def test_infers_fstring_dict_value():
5757
assert strings == ["Hello 123!"]
5858

5959

60-
def test_infers_string_format_value():
60+
def test_infers_string_format_value() -> None:
6161
source = """
6262
value = "abc"
6363
fstring = "Hello {0}!".format(value)
@@ -71,7 +71,7 @@ def test_infers_string_format_value():
7171
assert strings == ["Hello abc!"]
7272

7373

74-
def test_infers_fstring_values():
74+
def test_infers_fstring_values() -> None:
7575
source = """
7676
values_1 = ["abc", "def"]
7777
for value1 in values_1:
@@ -88,7 +88,7 @@ def test_infers_fstring_values():
8888
assert strings == ["Hello abc, ghi!", "Hello abc, jkl!", "Hello def, ghi!", "Hello def, jkl!"]
8989

9090

91-
def test_infers_externally_defined_value():
91+
def test_infers_externally_defined_value() -> None:
9292
state = CurrentSessionState()
9393
state.named_parameters = {"my-widget": "my-value"}
9494
source = """
@@ -103,7 +103,7 @@ def test_infers_externally_defined_value():
103103
assert strings == ["my-value"]
104104

105105

106-
def test_infers_externally_defined_values():
106+
def test_infers_externally_defined_values() -> None:
107107
state = CurrentSessionState()
108108
state.named_parameters = {"my-widget-1": "my-value-1", "my-widget-2": "my-value-2"}
109109
source = """
@@ -118,7 +118,7 @@ def test_infers_externally_defined_values():
118118
assert strings == ["my-value-1", "my-value-2"]
119119

120120

121-
def test_fails_to_infer_missing_externally_defined_value():
121+
def test_fails_to_infer_missing_externally_defined_value() -> None:
122122
state = CurrentSessionState()
123123
state.named_parameters = {"my-widget-1": "my-value-1", "my-widget-2": "my-value-2"}
124124
source = """
@@ -132,7 +132,7 @@ def test_fails_to_infer_missing_externally_defined_value():
132132
assert all(not value.is_inferred() for value in values)
133133

134134

135-
def test_survives_absence_of_externally_defined_values():
135+
def test_survives_absence_of_externally_defined_values() -> None:
136136
source = """
137137
name = "my-widget"
138138
value = dbutils.widgets.get(name)
@@ -144,7 +144,7 @@ def test_survives_absence_of_externally_defined_values():
144144
assert all(not value.is_inferred() for value in values)
145145

146146

147-
def test_infers_externally_defined_value_set():
147+
def test_infers_externally_defined_value_set() -> None:
148148
state = CurrentSessionState()
149149
state.named_parameters = {"my-widget": "my-value"}
150150
source = """

0 commit comments

Comments
 (0)