Making Type typechecks use class hierarchy.

cqfn · Jul 14, 2020 · 8076660 · 8076660
1 parent fac167b
commit 8076660
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/aibolit/ast_framework/_auxiliary_data.py b/aibolit/ast_framework/_auxiliary_data.py
@@ -23,6 +23,7 @@
 from typing import Dict, Set, Type, NamedTuple
 
 from javalang import tree
+from javalang.ast import Node
 
 from aibolit.ast_framework import ASTNodeType
 
@@ -31,7 +32,7 @@ class ASTNodeReference(NamedTuple):
     node_index: int
 
 
-javalang_to_ast_node_type: Dict[Type, ASTNodeType] = {
+javalang_to_ast_node_type: Dict[Type[Node], ASTNodeType] = {
     tree.Annotation: ASTNodeType.ANNOTATION,
     tree.AnnotationDeclaration: ASTNodeType.ANNOTATION_DECLARATION,
     tree.AnnotationMethod: ASTNodeType.ANNOTATION_METHOD,