fix

Author: asukaminato0721

crates/pyrefly_python/src/sys_info.rs

@@ -26,9 +26,11 @@ use ruff_python_ast::ExprAttribute;
 use ruff_python_ast::ExprBooleanLiteral;
 use ruff_python_ast::ExprCall;
 use ruff_python_ast::ExprList;
+use ruff_python_ast::ExprName;
 use ruff_python_ast::ExprNumberLiteral;
 use ruff_python_ast::ExprSet;
 use ruff_python_ast::ExprSlice;
+use ruff_python_ast::ExprStringLiteral;
 use ruff_python_ast::ExprSubscript;
 use ruff_python_ast::Stmt;
 use ruff_python_ast::StmtIf;
@@ -267,6 +269,14 @@ impl SysInfo {
     pub fn type_checking(&self) -> bool {
         self.0.key().type_checking
     }
+
+    pub fn with_platform(&self, platform: PythonPlatform) -> Self {
+        if self.type_checking() {
+            SysInfo::new(self.version(), platform)
+        } else {
+            SysInfo::new_without_type_checking(self.version(), platform)
+        }
+    }
 }
 
 impl<'de> Deserialize<'de> for SysInfo {
@@ -697,6 +707,83 @@ impl SysInfo {
     }
 }
 
+pub fn module_platform_guard(body: &[Stmt]) -> Option<PythonPlatform> {
+    body.iter().find_map(platform_guard_from_stmt)
+}
+
+fn platform_guard_from_stmt(stmt: &Stmt) -> Option<PythonPlatform> {
+    match stmt {
+        Stmt::Assert(x) => platform_guard_from_assert(&x.test),
+        Stmt::If(x) if x.elif_else_clauses.is_empty() && body_is_unconditional_raise(&x.body) => {
+            platform_guard_from_not_eq(&x.test)
+        }
+        _ => None,
+    }
+}
+
+fn body_is_unconditional_raise(body: &[Stmt]) -> bool {
+    matches!(body, [Stmt::Raise(_)])
+}
+
+fn platform_guard_from_assert(test: &Expr) -> Option<PythonPlatform> {
+    let (op, platform) = platform_compare(test)?;
+    match op {
+        CmpOp::Eq => Some(PythonPlatform::new(&platform)),
+        _ => None,
+    }
+}
+
+fn platform_guard_from_not_eq(test: &Expr) -> Option<PythonPlatform> {
+    let (op, platform) = platform_compare(test)?;
+    match op {
+        CmpOp::NotEq => Some(PythonPlatform::new(&platform)),
+        _ => None,
+    }
+}
+
+fn platform_compare(test: &Expr) -> Option<(CmpOp, String)> {
+    let Expr::Compare(compare) = test else {
+        return None;
+    };
+    if compare.ops.len() != 1 || compare.comparators.len() != 1 {
+        return None;
+    }
+    let op = compare.ops[0];
+    if !matches!(op, CmpOp::Eq | CmpOp::NotEq) {
+        return None;
+    }
+    let left = &compare.left;
+    let right = &compare.comparators[0];
+    let platform = extract_platform_literal(left, right)?;
+    Some((op, platform))
+}
+
+fn extract_platform_literal(left: &Expr, right: &Expr) -> Option<String> {
+    if is_sys_platform(left) {
+        string_literal(right)
+    } else if is_sys_platform(right) {
+        string_literal(left)
+    } else {
+        None
+    }
+}
+
+fn is_sys_platform(expr: &Expr) -> bool {
+    matches!(
+        expr,
+        Expr::Attribute(ExprAttribute { value, attr, .. })
+            if matches!(&**value, Expr::Name(ExprName { id, .. }) if id == "sys")
+                && attr.as_str() == "platform"
+    )
+}
+
+fn string_literal(expr: &Expr) -> Option<String> {
+    match expr {
+        Expr::StringLiteral(ExprStringLiteral { value, .. }) => Some(value.to_str().to_owned()),
+        _ => None,
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

pyrefly/lib/binding/bindings.rs

@@ -20,6 +20,7 @@ use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::nesting_context::NestingContext;
 use pyrefly_python::short_identifier::ShortIdentifier;
 use pyrefly_python::sys_info::SysInfo;
+use pyrefly_python::sys_info::module_platform_guard;
 use pyrefly_types::callable::FuncDefIndex;
 use pyrefly_types::type_alias::TypeAliasIndex;
 use pyrefly_types::type_info::JoinStyle;
@@ -473,6 +474,9 @@ impl Bindings {
         enable_trace: bool,
         untyped_def_behavior: UntypedDefBehavior,
     ) -> Self {
+        let override_sys_info =
+            module_platform_guard(&x.body).map(|platform| sys_info.with_platform(platform));
+        let sys_info = override_sys_info.as_ref().unwrap_or(sys_info);
         let mut builder = BindingsBuilder {
             module_info: module_info.dupe(),
             lookup,

pyrefly/lib/state/state.rs

@@ -34,6 +34,7 @@ use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModulePath;
 use pyrefly_python::module_path::ModulePathDetails;
 use pyrefly_python::sys_info::SysInfo;
+use pyrefly_python::sys_info::module_platform_guard;
 use pyrefly_types::type_alias::TypeAliasIndex;
 use pyrefly_util::arc_id::ArcId;
 use pyrefly_util::events::CategorizedEvents;
@@ -389,6 +390,14 @@ struct ModuleDataMut {
     rdeps: Mutex<HashSet<Handle>>,
 }
 
+fn module_sys_info_override(
+    sys_info: &SysInfo,
+    ast: Option<&ruff_python_ast::ModModule>,
+) -> Option<SysInfo> {
+    let ast = ast?;
+    module_platform_guard(&ast.body).map(|platform| sys_info.with_platform(platform))
+}
+
 impl ModuleData {
     /// Make a copy of the data that can be mutated.
     fn clone_for_mutation(&self) -> ModuleDataMut {
@@ -1014,11 +1023,18 @@ impl<'a> Transaction<'a> {
             let require = guard.require();
             let stdlib = self.get_stdlib(&module_data.handle);
             let config = module_data.config.read();
+            let sys_info_override = module_sys_info_override(
+                module_data.handle.sys_info(),
+                module_data.state.get_ast().as_deref(),
+            );
+            let sys_info = sys_info_override
+                .as_ref()
+                .unwrap_or(module_data.handle.sys_info());
             let ctx = Context {
                 require,
                 module: module_data.handle.module(),
                 path: module_data.handle.path(),
-                sys_info: module_data.handle.sys_info(),
+                sys_info,
                 memory: &self.memory_lookup(),
                 uniques: &self.data.state.uniques,
                 stdlib: &stdlib,
@@ -1402,13 +1418,17 @@ impl<'a> Transaction<'a> {
             .dupe()
     }
 
-    pub fn get_stdlib(&self, handle: &Handle) -> Arc<Stdlib> {
+    pub fn get_stdlib_for_sys_info(&self, sys_info: &SysInfo) -> Arc<Stdlib> {
         if self.data.stdlib.len() == 1 {
             // Since we know our one must exist, we can shortcut
             return self.data.stdlib.first().unwrap().1.dupe();
         }
 
-        self.data.stdlib.get(handle.sys_info()).unwrap().dupe()
+        self.data.stdlib.get(sys_info).unwrap().dupe()
+    }
+
+    pub fn get_stdlib(&self, handle: &Handle) -> Arc<Stdlib> {
+        self.get_stdlib_for_sys_info(handle.sys_info())
     }
 
     fn compute_stdlib(&mut self, sys_infos: SmallSet<SysInfo>) {
@@ -1983,14 +2003,17 @@ impl<'a> TransactionHandle<'a> {
         path: Option<&ModulePath>,
         dep: ModuleDep,
     ) -> FindingOrError<ArcId<ModuleDataMut>> {
+        let sys_info_override = module_sys_info_override(
+            self.module_data.handle.sys_info(),
+            self.module_data.state.get_ast().as_deref(),
+        );
+        let sys_info = sys_info_override
+            .as_ref()
+            .unwrap_or(self.module_data.handle.sys_info());
         let handle = match path {
             Some(path) => {
                 // Explicit path — already resolved. Bypass imports entirely.
-                FindingOrError::new_finding(Handle::new(
-                    module,
-                    path.dupe(),
-                    self.module_data.handle.sys_info().dupe(),
-                ))
+                FindingOrError::new_finding(Handle::new(module, path.dupe(), sys_info.dupe()))
             }
             None => {
                 // No path — needs find_import. Check imports cache first.
@@ -2010,13 +2033,7 @@ impl<'a> TransactionHandle<'a> {
                         finding
                     }
                 };
-                path.map(|path| {
-                    Handle::new(
-                        module,
-                        path.dupe(),
-                        self.module_data.handle.sys_info().dupe(),
-                    )
-                })
+                path.map(|path| Handle::new(module, path.dupe(), sys_info.dupe()))
             }
         };
 

pyrefly/lib/test/sys_info.rs

@@ -204,6 +204,22 @@ assert_type(Z(), str)
 "#,
 );
 
+testcase!(
+    test_module_platform_guard_imports,
+    TestEnv::one(
+        "winonly",
+        r#"
+import sys
+assert sys.platform == "win32"
+import winreg
+x = winreg.OpenKey
+"#,
+    ),
+    r#"
+import winonly
+"#,
+);
+
 testcase!(
     test_os_name,
     r#"