SONARPY-945 Use precomputed Typeshed symbols for custom stub files

Author: guillaume-dequenne

python-frontend/src/main/java/org/sonar/python/types/TypeShed.java

@@ -58,7 +58,6 @@
 import static org.sonar.plugins.python.api.types.BuiltinTypes.NONE_TYPE;
 import static org.sonar.plugins.python.api.types.BuiltinTypes.STR;
 import static org.sonar.plugins.python.api.types.BuiltinTypes.TUPLE;
-import static org.sonar.python.types.TypeShedThirdParties.getModuleSymbols;
 
 public class TypeShed {
 
@@ -71,8 +70,9 @@ public class TypeShed {
   private static final String THIRD_PARTY_2 = "typeshed/third_party/2/";
   private static final String THIRD_PARTY_3 = "typeshed/third_party/3/";
   private static final String CUSTOM_THIRD_PARTY = "custom/";
-  private static final String PROTOBUF = "protobuf/";
-  private static final String PROTOBUF_THIRD_PARTY = "protobuf/stubs/";
+  private static final String PROTOBUF_CUSTOM_STUBS = "custom_protobuf/";
+  private static final String PROTOBUF = "stdlib_protobuf/";
+  private static final String PROTOBUF_THIRD_PARTY = "third_party_protobuf/";
   private static final String BUILTINS_FQN = "builtins";
   private static final String BUILTINS_PREFIX = BUILTINS_FQN + ".";
   // Those fundamentals builtins symbols need not to be ambiguous for the frontend to work properly
@@ -96,7 +96,7 @@ private TypeShed() {
   public static Map<String, Symbol> builtinSymbols() {
     if ((TypeShed.builtins == null)) {
       supportedPythonVersions = ProjectPythonVersion.currentVersions().stream().map(PythonVersionUtils.Version::serializedValue).collect(Collectors.toSet());
-      Map<String, Symbol> builtins = getSymbolsFromProtobufModule(BUILTINS_FQN, false);
+      Map<String, Symbol> builtins = getSymbolsFromProtobufModule(BUILTINS_FQN, PROTOBUF);
       builtins.put(NONE_TYPE, new ClassSymbolImpl(NONE_TYPE, NONE_TYPE));
       TypeShed.builtins = Collections.unmodifiableMap(builtins);
       TypeShed.builtinGlobalSymbols.put("", new HashSet<>(builtins.values()));
@@ -266,18 +266,18 @@ private static Map<String, Symbol> searchTypeShedForModule(String moduleName) {
       return new HashMap<>();
     }
     modulesInProgress.add(moduleName);
-    Map<String, Symbol> customSymbols = getModuleSymbols(moduleName, CUSTOM_THIRD_PARTY, builtinGlobalSymbols);
+    Map<String, Symbol> customSymbols = getSymbolsFromProtobufModule(moduleName, PROTOBUF_CUSTOM_STUBS);
     if (!customSymbols.isEmpty()) {
       modulesInProgress.remove(moduleName);
       return customSymbols;
     }
-    Map<String, Symbol> symbolsFromProtobuf = getSymbolsFromProtobufModule(moduleName, false);
+    Map<String, Symbol> symbolsFromProtobuf = getSymbolsFromProtobufModule(moduleName, PROTOBUF);
     if (!symbolsFromProtobuf.isEmpty()) {
       modulesInProgress.remove(moduleName);
       return symbolsFromProtobuf;
     }
 
-    Map<String, Symbol> thirdPartySymbols = getSymbolsFromProtobufModule(moduleName, true);
+    Map<String, Symbol> thirdPartySymbols = getSymbolsFromProtobufModule(moduleName, PROTOBUF_THIRD_PARTY);
     modulesInProgress.remove(moduleName);
     return thirdPartySymbols;
   }
@@ -308,9 +308,8 @@ private static boolean isAmbiguousSymbolOfClasses(Symbol symbol) {
     return false;
   }
 
-  private static Map<String, Symbol> getSymbolsFromProtobufModule(String moduleName, boolean isThirdParty) {
-    String protobufDir = isThirdParty ? PROTOBUF_THIRD_PARTY : PROTOBUF;
-    InputStream resource = TypeShed.class.getResourceAsStream(protobufDir + moduleName + ".protobuf");
+  private static Map<String, Symbol> getSymbolsFromProtobufModule(String moduleName, String dirName) {
+    InputStream resource = TypeShed.class.getResourceAsStream(dirName + moduleName + ".protobuf");
     if (resource == null) {
       return Collections.emptyMap();
     }

python-frontend/src/test/java/org/sonar/python/types/CustomStubSanityTest.java

@@ -264,6 +264,15 @@ public void package_django_class_property_type() {
     assertThat(((ClassSymbol) requestSymbol).declaredMembers().iterator().next().annotatedTypeName()).isEqualTo("django.http.request.QueryDict");
   }
 
+  @Test
+  public void package_lxml_reexported_symbol_fqn() {
+    Map<String, Symbol> lxmlEtreeSymbols = symbolsForModule("lxml.etree");
+    Symbol elementTreeSymbol = lxmlEtreeSymbols.get("ElementTree");
+    assertThat(elementTreeSymbol.kind()).isEqualTo(Kind.CLASS);
+    // FIXME: Original FQN is "xml.etree.ElementTree.ElementTree" and we should be able to retrieve it somehow
+    assertThat(elementTreeSymbol.fullyQualifiedName()).isEqualTo("lxml.etree.ElementTree");
+  }
+
   @Test
   public void package_sqlite3_connect_type_in_ambiguous_symbol() {
     Map<String, Symbol> djangoSymbols = symbolsForModule("sqlite3");

python-frontend/src/test/java/org/sonar/python/types/TypeShedTest.java

@@ -33,6 +33,7 @@
 logger = logging.getLogger(__name__)
 
 DEFAULT_EXPORTED_VARS = ["__name__", "__doc__", "__file__", "__package__"]
+SONAR_CUSTOM_BASE_CLASS = "SonarPythonAnalyzerFakeStub.CustomStubBase"
 
 
 class ParamKind(Enum):
@@ -337,7 +338,8 @@ def __init__(self, name: str, fullname: str, is_imported_module=False, type_desc
 
     @classmethod
     def from_var(cls, var: mpn.Var, name: str = None):
-        return cls(var.name if name is None else name, var.fullname, type_descriptor=TypeDescriptor(var.type) if var.type else None)
+        return cls(var.name if name is None else name, var.fullname,
+                   type_descriptor=TypeDescriptor(var.type) if var.type else None)
 
     def __eq__(self, other):
         return isinstance(other, VarSymbol) and self.to_proto() == other.to_proto()
@@ -361,6 +363,9 @@ def __init__(self, mypy_file: mpn.MypyFile):
         self.vars = []
         for key in mypy_file.names:
             name = mypy_file.names.get(key)
+            if name.fullname == SONAR_CUSTOM_BASE_CLASS:
+                # Ignore custom stub name
+                continue
             symbol_table_node = name.node
             if isinstance(symbol_table_node, mpn.FuncDef):
                 self.functions.append(FunctionSymbol(symbol_table_node, name=key))
@@ -528,10 +533,10 @@ def extract_return_type(func_def: mpn.FuncDef):
     return TypeDescriptor(func_type.ret_type)
 
 
-def save_module(ms: Union[ModuleSymbol, MergedModuleSymbol], is_debug=False, debug_dir="output", is_stdlib=True):
+def save_module(ms: Union[ModuleSymbol, MergedModuleSymbol], dir_name="stdlib_protobuf",
+                is_debug=False, debug_dir="output"):
     ms_pb = ms.to_proto()
-    save_dir = "../../src/main/resources/org/sonar/python/types/protobuf" if not is_debug else f"../{debug_dir}"
-    save_dir = save_dir if is_stdlib else os.path.join(save_dir, "stubs")
+    save_dir = f"../../src/main/resources/org/sonar/python/types/{dir_name}" if not is_debug else f"../{debug_dir}"
     save_string = ms_pb.SerializeToString() if not is_debug else str(ms_pb)
     open_mode = "wb" if not is_debug else "w"
     save_dir_path = os.path.join(CURRENT_PATH, save_dir)

python-frontend/typeshed_serializer/serializer/symbols.py

@@ -28,6 +28,8 @@
 STUBS_PATH = "../resources/typeshed/stubs"
 CURRENT_PATH = os.path.dirname(__file__)
 THIRD_PARTIES_STUBS = os.listdir(os.path.join(CURRENT_PATH, STUBS_PATH))
+CUSTOM_STUBS_PATH = "../resources/custom"
+SONAR_CUSTOM_BASE_STUB_MODULE = "SonarPythonAnalyzerFakeStub"
 
 
 def get_options(python_version=(3, 8)):
@@ -83,6 +85,12 @@ def walk_typeshed_third_parties(opt: options.Options = get_options()):
     return build_result, source_paths
 
 
+def walk_custom_stubs(opt: options.Options = get_options()):
+    source_list, source_paths = get_sources(CUSTOM_STUBS_PATH, False)
+    build_result = build.build(source_list, opt)
+    return build_result, source_paths
+
+
 def get_sources(relative_path: str, generate_python2: bool):
     source_list = []
     source_paths = set()
@@ -118,7 +126,21 @@ def serialize_typeshed_stdlib(output_dir_name="output", python_version=(3, 8), i
     build_result, _ = walk_typeshed_stdlib(opt)
     for file in build_result.files:
         module_symbol = symbols.ModuleSymbol(build_result.files.get(file))
-        symbols.save_module(module_symbol, is_debug=is_debug, debug_dir=output_dir_name)
+        symbols.save_module(module_symbol, "stdlib_protobuf", is_debug=is_debug, debug_dir=output_dir_name)
+
+
+def serialize_custom_stubs(output_dir_name="output", python_version=(3, 8), is_debug=False):
+    path = os.path.join(CURRENT_PATH, CUSTOM_STUBS_PATH)
+    opt = get_options(python_version)
+    build_result, _ = walk_custom_stubs(opt)
+    for file in build_result.files:
+        if file == SONAR_CUSTOM_BASE_STUB_MODULE:
+            continue
+        current_file = build_result.files.get(file)
+        if not current_file.path.startswith(path):
+            continue
+        module_symbol = symbols.ModuleSymbol(current_file)
+        symbols.save_module(module_symbol, "custom_protobuf", is_debug=is_debug, debug_dir=output_dir_name)
 
 
 def serialize_typeshed_stdlib_multiple_python_version():
@@ -129,14 +151,16 @@ def serialize_typeshed_stdlib_multiple_python_version():
 
 
 def save_merged_symbols(is_debug=False, is_third_parties=False):
+    dir_name = "third_party_protobuf" if is_third_parties else "stdlib_protobuf"
     merged_modules = symbols_merger.merge_multiple_python_versions(is_third_parties)
     for mod in merged_modules:
-        symbols.save_module(merged_modules[mod], is_debug=is_debug, debug_dir="output_merge", is_stdlib=not is_third_parties)
+        symbols.save_module(merged_modules[mod], dir_name, is_debug=is_debug, debug_dir="output_merge")
 
 
 def main():
     save_merged_symbols()
     save_merged_symbols(is_third_parties=True)
+    serialize_custom_stubs()
 
 
 if __name__ == '__main__':

python-frontend/typeshed_serializer/serializer/typeshed_serializer.py

@@ -33,6 +33,13 @@ def typeshed_stdlib():
     return build_result
 
 
+@pytest.fixture(scope="session")
+def typeshed_custom_stubs():
+    build_result, _ = typeshed_serializer.walk_custom_stubs()
+    assert len(build_result.errors) == 0
+    return build_result
+
+
 @pytest.fixture(scope="session")
 def fake_module_36_38():
     fake_module_path = os.path.join(os.path.dirname(__file__), "resources/fakemodule.pyi")

python-frontend/typeshed_serializer/tests/conftest.py

@@ -36,6 +36,15 @@ def test_serialize_typeshed_stdlib(typeshed_stdlib):
     assert symbols.save_module.call_count == len(typeshed_stdlib.files)
 
 
+def test_serialize_custom_stubs(typeshed_custom_stubs):
+    typeshed_serializer.walk_custom_stubs = Mock(return_value=(typeshed_custom_stubs, set()))
+    symbols.save_module = Mock()
+    typeshed_serializer.serialize_custom_stubs()
+    assert typeshed_serializer.walk_custom_stubs.call_count == 1
+    # Not every files from "typeshed_custom_stubs" build are serialized, as some are builtins
+    assert symbols.save_module.call_count == 50
+
+
 def test_all_third_parties_are_serialized(typeshed_third_parties):
     stub_modules = set()
     for stub_folder in typeshed_serializer.THIRD_PARTIES_STUBS: