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`

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

@@ -66,15 +66,15 @@ def _infer_values(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
             yield from cls._safe_infer_internal(node)
 
     @classmethod
-    def _safe_infer_internal(cls, node: NodeNG):
+    def _safe_infer_internal(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
         try:
             yield from cls._unsafe_infer_internal(node)
-        except InferenceError as e:
-            logger.debug(f"When inferring {node}", exc_info=e)
+        except (InferenceError, KeyError) as e:
+            logger.debug(f"When inferring: {node}", exc_info=e)
             yield [Uninferable]
 
     @classmethod
-    def _unsafe_infer_internal(cls, node: NodeNG):
+    def _unsafe_infer_internal(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
         all_inferred = node.inferred()
         if len(all_inferred) == 0 and isinstance(node, Instance):
             yield [node]
@@ -148,7 +148,7 @@ def infer_call_result(self, context: InferenceContext | None = None, **_):  # ca
 class _LocalInferredValue(InferredValue):
 
     @classmethod
-    def do_infer_values(cls, node: NodeNG):
+    def do_infer_values(cls, node: NodeNG) -> Iterator[Iterable[NodeNG]]:
         yield from cls._infer_values(node)
 
 

tests/unit/source_code/python/test_python_infer.py

@@ -5,31 +5,31 @@
 from databricks.labs.ucx.source_code.python.python_infer import InferredValue
 
 
-def test_infers_empty_list():
+def test_infers_empty_list() -> None:
     tree = Tree.parse("a=[]")
     nodes = tree.locate(Assign, [])
     tree = Tree(nodes[0].value)
     values = list(InferredValue.infer_from_node(tree.node))
     assert not values
 
 
-def test_infers_empty_tuple():
+def test_infers_empty_tuple() -> None:
     tree = Tree.parse("a=tuple()")
     nodes = tree.locate(Assign, [])
     tree = Tree(nodes[0].value)
     values = list(InferredValue.infer_from_node(tree.node))
     assert not values
 
 
-def test_infers_empty_set():
+def test_infers_empty_set() -> None:
     tree = Tree.parse("a={}")
     nodes = tree.locate(Assign, [])
     tree = Tree(nodes[0].value)
     values = list(InferredValue.infer_from_node(tree.node))
     assert not values
 
 
-def test_infers_fstring_value():
+def test_infers_fstring_value() -> None:
     source = """
 value = "abc"
 fstring = f"Hello {value}!"
@@ -43,7 +43,7 @@ def test_infers_fstring_value():
     assert strings == ["Hello abc!"]
 
 
-def test_infers_fstring_dict_value():
+def test_infers_fstring_dict_value() -> None:
     source = """
 value = { "abc": 123 }
 fstring = f"Hello {value['abc']}!"
@@ -57,7 +57,7 @@ def test_infers_fstring_dict_value():
     assert strings == ["Hello 123!"]
 
 
-def test_infers_string_format_value():
+def test_infers_string_format_value() -> None:
     source = """
 value = "abc"
 fstring = "Hello {0}!".format(value)
@@ -71,7 +71,7 @@ def test_infers_string_format_value():
     assert strings == ["Hello abc!"]
 
 
-def test_infers_fstring_values():
+def test_infers_fstring_values() -> None:
     source = """
 values_1 = ["abc", "def"]
 for value1 in values_1:
@@ -88,7 +88,7 @@ def test_infers_fstring_values():
     assert strings == ["Hello abc, ghi!", "Hello abc, jkl!", "Hello def, ghi!", "Hello def, jkl!"]
 
 
-def test_infers_externally_defined_value():
+def test_infers_externally_defined_value() -> None:
     state = CurrentSessionState()
     state.named_parameters = {"my-widget": "my-value"}
     source = """
@@ -103,7 +103,7 @@ def test_infers_externally_defined_value():
     assert strings == ["my-value"]
 
 
-def test_infers_externally_defined_values():
+def test_infers_externally_defined_values() -> None:
     state = CurrentSessionState()
     state.named_parameters = {"my-widget-1": "my-value-1", "my-widget-2": "my-value-2"}
     source = """
@@ -118,7 +118,7 @@ def test_infers_externally_defined_values():
     assert strings == ["my-value-1", "my-value-2"]
 
 
-def test_fails_to_infer_missing_externally_defined_value():
+def test_fails_to_infer_missing_externally_defined_value() -> None:
     state = CurrentSessionState()
     state.named_parameters = {"my-widget-1": "my-value-1", "my-widget-2": "my-value-2"}
     source = """
@@ -132,7 +132,7 @@ def test_fails_to_infer_missing_externally_defined_value():
     assert all(not value.is_inferred() for value in values)
 
 
-def test_survives_absence_of_externally_defined_values():
+def test_survives_absence_of_externally_defined_values() -> None:
     source = """
     name = "my-widget"
     value = dbutils.widgets.get(name)
@@ -144,7 +144,7 @@ def test_survives_absence_of_externally_defined_values():
     assert all(not value.is_inferred() for value in values)
 
 
-def test_infers_externally_defined_value_set():
+def test_infers_externally_defined_value_set() -> None:
     state = CurrentSessionState()
     state.named_parameters = {"my-widget": "my-value"}
     source = """