Fix incorrect method type map resolution order

Author: vemel

mypy_boto3_builder/type_maps/method_type_map.py

@@ -98,41 +98,6 @@
 }
 
 
-def _get_from_method_map(
-    method_name: str,
-    argument_name: str,
-    method_type_map: MethodTypeMap,
-) -> FakeAnnotation | None:
-    if ALL in method_type_map:
-        operation_type_map = method_type_map[ALL]
-        if argument_name in operation_type_map:
-            return operation_type_map[argument_name]
-
-    if method_name in method_type_map:
-        operation_type_map = method_type_map[method_name]
-        if argument_name in operation_type_map:
-            return operation_type_map[argument_name]
-
-    return None
-
-
-def _get_from_class_map(
-    class_name: str,
-    method_name: str,
-    argument_name: str,
-    class_type_map: ClassTypeMap,
-) -> FakeAnnotation | None:
-    if class_name in class_type_map:
-        result = _get_from_method_map(method_name, argument_name, class_type_map[class_name])
-        if result:
-            return result
-    if ALL in class_type_map:
-        result = _get_from_method_map(method_name, argument_name, class_type_map[ALL])
-        if result:
-            return result
-    return None
-
-
 def _get_from_service_map(
     service_name: ServiceName,
     class_name: str,
@@ -143,12 +108,25 @@ def _get_from_service_map(
     if service_name not in service_type_map:
         return None
 
-    return _get_from_class_map(
-        class_name,
-        method_name,
-        argument_name,
-        service_type_map[service_name],
+    checks = (
+        (class_name, method_name, argument_name),
+        (class_name, ALL, argument_name),
+        (ALL, method_name, argument_name),
+        (ALL, ALL, argument_name),
     )
+    class_type_map = service_type_map[service_name]
+
+    for class_name, method_name, argument_name in checks:
+        if class_name not in class_type_map:
+            continue
+
+        method_type_map = class_type_map[class_name]
+
+        if method_name not in method_type_map:
+            continue
+
+        operation_type_map = method_type_map[method_name]
+        return operation_type_map.get(argument_name)
 
 
 def get_method_type_stub(

tests/type_maps/test_method_type_map.py

@@ -15,6 +15,15 @@ def test_get_method_type_stub(self) -> None:
             get_method_type_stub(ServiceNameCatalog.logs, "Client", "copy_object", "Unknown")
             is None
         )
+        assert get_method_type_stub(
+            ServiceNameCatalog.sqs, "SomeResource", "receive_messages", "AttributeNames"
+        )
+        assert (
+            get_method_type_stub(
+                ServiceNameCatalog.sqs, "SomeResource", "receive_messages", "Attribute"
+            )
+            is None
+        )
 
     def test_get_default_value_stub(self) -> None:
         assert get_default_value_stub(ServiceNameCatalog.glacier, "Client", "*", "accountId")