Detect partial stub packages (#13)

Work towards #11

This is a `semanal_pass1` feature that would be good to access without
full deserialization, since it is needed early during graph loading.

Author: ilevkivskyi

src/lib.rs

@@ -47,7 +47,7 @@ fn parse(
     platform: Option<String>,
     always_true: Option<Vec<String>>,
     always_false: Option<Vec<String>>,
-) -> PyResult<(Vec<u8>, Vec<PyObject>, Vec<PyObject>, Vec<u8>)> {
+) -> PyResult<(Vec<u8>, Vec<PyObject>, Vec<PyObject>, Vec<u8>, bool)> {
     // Get defaults from Python if not provided
     let python_version = match python_version {
         Some(v) => v,
@@ -63,7 +63,7 @@ fn parse(
     let always_false = always_false.unwrap_or_default();
 
     let path = Path::new(&fnam);
-    let (ast_bytes, syntax_errors, type_ignore_lines, import_bytes) = py
+    let (ast_bytes, syntax_errors, type_ignore_lines, import_bytes, is_partial_package) = py
         .allow_threads(|| {
             serialize_ast::serialize_python_file(
                 path,
@@ -102,7 +102,7 @@ fn parse(
         .collect();
     let py_type_ignores = py_type_ignores?;
 
-    Ok((ast_bytes, py_errors, py_type_ignores, import_bytes))
+    Ok((ast_bytes, py_errors, py_type_ignores, import_bytes, is_partial_package))
 }
 
 /// Get the default Python version from sys.version_info

src/serialize_ast.rs

@@ -171,10 +171,14 @@ pub(crate) fn serialize_python_file(
     file_path: &Path,
     skip_function_bodies: bool,
     options: Options,
-) -> Result<(Vec<u8>, Vec<SyntaxError>, Vec<(usize, Vec<String>)>, Vec<u8>)> {
+) -> Result<(Vec<u8>, Vec<SyntaxError>, Vec<(usize, Vec<String>)>, Vec<u8>, bool)> {
     let source_type = PySourceType::from(file_path);
     let source_text = std::fs::read_to_string(file_path)?;
     let line_index = LineIndex::from_source_text(&source_text);
+    let is_stub_package = match file_path.file_name() {
+        Some(file) => file.as_encoded_bytes() == b"__init__.pyi",
+        _ => false,
+    };
 
     // Check if file is all ASCII and build per-line non-ASCII flags if needed
     let is_all_ascii = source_text.is_ascii();
@@ -221,6 +225,7 @@ pub(crate) fn serialize_python_file(
         options,
         current_unreachable: false,
         current_mypy_only: false,
+        top_level_getattr: false,
     };
     parsed.syntax().serialize(&mut ser);
 
@@ -233,7 +238,10 @@ pub(crate) fn serialize_python_file(
         Some(ser.lines_with_non_ascii),
     );
 
-    Ok((ser.bytes, syntax_errors, type_ignore_lines, import_bytes))
+    // Return this directly to caller, so that it can check this without deserialization
+    let is_partial_package = is_stub_package && ser.top_level_getattr;
+
+    Ok((ser.bytes, syntax_errors, type_ignore_lines, import_bytes, is_partial_package))
 }
 
 // Bit flags for import statement metadata
@@ -279,6 +287,7 @@ struct Serializer<'a> {
     options: Options,           // Reachability analysis options
     current_unreachable: bool,  // Whether we're currently in an unreachable block
     current_mypy_only: bool,    // Whether we're currently in a mypy-only block (e.g., if TYPE_CHECKING)
+    top_level_getattr: bool,    // Does module have top-level __getattr__() function
 }
 
 impl<'a> Serializer<'a> {
@@ -864,6 +873,10 @@ impl Ser for ast::Stmt {
                     true
                 };
 
+                if !ser.in_class && !ser.in_function && f.name.as_str() == "__getattr__" {
+                    ser.top_level_getattr = true;
+                };
+
                 // Body - mark that we're inside a function
                 let was_in_function = ser.in_function;
                 ser.in_function = true;
@@ -2580,6 +2593,7 @@ pub fn serialize_imports(
         options: Options::default(),
         current_unreachable: false,
         current_mypy_only: false,
+        top_level_getattr: false,
     };
 
     // Write list of imports
@@ -2708,6 +2722,7 @@ mod tests {
             options: Options::default(),
             current_unreachable: false,
             current_mypy_only: false,
+            top_level_getattr: false,
         }
     }