From e627fee392077dcb25c2c502322bc9b7e9473bcf Mon Sep 17 00:00:00 2001
From: Cor <jczuurmond@protonmail.com>
Date: Thu, 10 Oct 2024 16:20:31 +0200
Subject: [PATCH] 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`
---
 .../ucx/source_code/python/python_infer.py    | 10 ++++----
 .../source_code/python/test_python_infer.py   | 24 +++++++++----------
 2 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/src/databricks/labs/ucx/source_code/python/python_infer.py b/src/databricks/labs/ucx/source_code/python/python_infer.py
index 6727745280..9eda2d0556 100644
--- a/src/databricks/labs/ucx/source_code/python/python_infer.py
+++ b/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)
 
 
diff --git a/tests/unit/source_code/python/test_python_infer.py b/tests/unit/source_code/python/test_python_infer.py
index 9afc690a7d..a0ffa6fe72 100644
--- a/tests/unit/source_code/python/test_python_infer.py
+++ b/tests/unit/source_code/python/test_python_infer.py
@@ -5,7 +5,7 @@
 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)
@@ -13,7 +13,7 @@ def test_infers_empty_list():
     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)
@@ -21,7 +21,7 @@ def test_infers_empty_tuple():
     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)
@@ -29,7 +29,7 @@ def test_infers_empty_set():
     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 = """