diff --git a/src/protocol.rs b/src/protocol.rs index f140171..9374533 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -321,6 +321,12 @@ pub enum TypeDescriptor { #[serde(skip_serializing_if = "Option::is_none")] display: Option, name: String, + #[serde(skip_serializing_if = "Option::is_none")] + variance: Option, + #[serde(skip_serializing_if = "Option::is_none")] + upper_bound: Option, + #[serde(skip_serializing_if = "Vec::is_empty")] + constraints: Vec, }, #[serde(rename_all = "camelCase")] diff --git a/src/registry.rs b/src/registry.rs index ea70d85..2d9c5e5 100644 --- a/src/registry.rs +++ b/src/registry.rs @@ -3,6 +3,7 @@ use ty_python_semantic::Db; use ty_python_semantic::types::list_members; use ty_python_semantic::types::{ ClassLiteral, GenericContext, LiteralValueTypeKind, ParameterKind, Type, TypeGuardLike, + TypeVarVariance, }; use crate::protocol::{ClassMemberInfo, ParameterInfo, TypeDescriptor, TypeId, TypedDictFieldInfo}; @@ -506,7 +507,33 @@ impl<'db> TypeRegistry<'db> { Type::TypeVar(bound_tv) => { let display = self.display_string(ty, db); let name = bound_tv.name(db).to_string(); - TypeDescriptor::TypeVar { display, name } + let typevar = bound_tv.typevar(db); + + let variance = Some( + match bound_tv.variance(db) { + TypeVarVariance::Covariant => "covariant", + TypeVarVariance::Contravariant => "contravariant", + TypeVarVariance::Invariant | TypeVarVariance::Bivariant => "invariant", + } + .to_string(), + ); + + let upper_bound = typevar + .upper_bound(db) + .map(|bound| self.register_component(bound, db)); + + let constraints: Vec<_> = typevar + .constraints(db) + .map(|cs| cs.iter().map(|&c| self.register_component(c, db)).collect()) + .unwrap_or_default(); + + TypeDescriptor::TypeVar { + display, + name, + variance, + upper_bound, + constraints, + } } Type::TypeAlias(_) => { diff --git a/tests/integration/main.rs b/tests/integration/main.rs index eab06fe..2591141 100644 --- a/tests/integration/main.rs +++ b/tests/integration/main.rs @@ -662,3 +662,198 @@ fn test_module_names() { instance_type.get("moduleName") ); } + +#[test] +fn test_typevar_variance_covariant() { + let dir = create_test_project(&[( + "v.py", + "from typing import TypeVar\nT_co = TypeVar('T_co', covariant=True)\nclass Box[T_co]: ...\n", + )]); + + let responses = run_session(&[ + &initialize_request(dir.path().to_str().unwrap(), 1), + &get_types_request("v.py", 2), + &shutdown_request(99), + ]); + + let result = &responses[1]["result"]; + let types: TypeMap = serde_json::from_value(result["types"].clone()).unwrap(); + + // Find a TypeVar — any T_co should have a variance field + let tv = types + .values() + .find(|t| t["kind"] == "typeVar" && t["name"] == "T_co"); + if let Some(tv) = tv { + assert!( + tv.get("variance").is_some(), + "typeVar should have a variance field, got {:?}", + tv + ); + } +} + +#[test] +fn test_typevar_variance_on_pep695() { + let dir = create_test_project(&[("vp.py", "def identity[T](x: T) -> T: return x\n")]); + + let responses = run_session(&[ + &initialize_request(dir.path().to_str().unwrap(), 1), + &get_types_request("vp.py", 2), + &shutdown_request(99), + ]); + + let result = &responses[1]["result"]; + let types: TypeMap = serde_json::from_value(result["types"].clone()).unwrap(); + + let tv = types + .values() + .find(|t| t["kind"] == "typeVar" && t["name"] == "T") + .expect("should have a TypeVar T"); + + // PEP 695 typevars should have an inferred variance + assert!( + tv.get("variance").is_some(), + "PEP 695 typeVar should have variance, got {:?}", + tv + ); + let variance = tv["variance"].as_str().unwrap(); + assert!( + ["covariant", "contravariant", "invariant"].contains(&variance), + "variance should be a valid value, got {:?}", + variance + ); +} + +#[test] +fn test_typevar_with_upper_bound() { + let dir = create_test_project(&[( + "b.py", + "from typing import TypeVar\nT = TypeVar('T', bound=int)\ndef f(x: T) -> T: return x\n", + )]); + + let responses = run_session(&[ + &initialize_request(dir.path().to_str().unwrap(), 1), + &get_types_request("b.py", 2), + &shutdown_request(99), + ]); + + let result = &responses[1]["result"]; + let types: TypeMap = serde_json::from_value(result["types"].clone()).unwrap(); + + // Find the TypeVar T used as a type parameter of function f + let func_type = types + .values() + .find(|t| t["kind"] == "function" && t["name"] == "f") + .expect("should have function f"); + let type_params = func_type["typeParameters"] + .as_array() + .expect("should have typeParameters"); + assert_eq!(type_params.len(), 1); + + let tv_id = type_params[0].to_string(); + let tv = &types[&tv_id]; + assert_eq!(tv["kind"], "typeVar"); + assert_eq!(tv["name"], "T"); + + // Should have an upperBound pointing to int + assert!( + tv.get("upperBound").is_some(), + "bounded TypeVar should have upperBound, got {:?}", + tv + ); + let bound_id = tv["upperBound"].to_string(); + let bound_type = &types[&bound_id]; + assert_eq!( + bound_type["className"], "int", + "upper bound should be int, got {:?}", + bound_type + ); +} + +#[test] +fn test_typevar_with_constraints() { + let dir = create_test_project(&[( + "c.py", + "from typing import TypeVar\nT = TypeVar('T', int, str)\ndef f(x: T) -> T: return x\n", + )]); + + let responses = run_session(&[ + &initialize_request(dir.path().to_str().unwrap(), 1), + &get_types_request("c.py", 2), + &shutdown_request(99), + ]); + + let result = &responses[1]["result"]; + let types: TypeMap = serde_json::from_value(result["types"].clone()).unwrap(); + + // Find the TypeVar T used as a type parameter of function f + let func_type = types + .values() + .find(|t| t["kind"] == "function" && t["name"] == "f") + .expect("should have function f"); + let type_params = func_type["typeParameters"] + .as_array() + .expect("should have typeParameters"); + assert_eq!(type_params.len(), 1); + + let tv_id = type_params[0].to_string(); + let tv = &types[&tv_id]; + assert_eq!(tv["kind"], "typeVar"); + assert_eq!(tv["name"], "T"); + + // Should have constraints with two entries (int and str) + let constraints = tv["constraints"] + .as_array() + .expect("constrained TypeVar should have constraints array"); + assert_eq!( + constraints.len(), + 2, + "TypeVar(T, int, str) should have 2 constraints, got {:?}", + constraints + ); + + // Verify the constraint types are int and str + let constraint_names: Vec<&str> = constraints + .iter() + .filter_map(|c| { + let cid = c.to_string(); + types[&cid]["className"].as_str() + }) + .collect(); + assert!( + constraint_names.contains(&"int") && constraint_names.contains(&"str"), + "constraints should include int and str, got {:?}", + constraint_names + ); +} + +#[test] +fn test_typevar_no_bounds_no_constraints() { + let dir = create_test_project(&[("nb.py", "def identity[T](x: T) -> T: return x\n")]); + + let responses = run_session(&[ + &initialize_request(dir.path().to_str().unwrap(), 1), + &get_types_request("nb.py", 2), + &shutdown_request(99), + ]); + + let result = &responses[1]["result"]; + let types: TypeMap = serde_json::from_value(result["types"].clone()).unwrap(); + + let tv = types + .values() + .find(|t| t["kind"] == "typeVar" && t["name"] == "T") + .expect("should have TypeVar T"); + + // Unbounded TypeVar should not have upperBound or constraints keys + assert!( + tv.get("upperBound").is_none(), + "unbounded TypeVar should not have upperBound, got {:?}", + tv + ); + assert!( + tv.get("constraints").is_none(), + "unconstrained TypeVar should not have constraints key, got {:?}", + tv + ); +}