Fix hover type for all definition contexts

Summary:
Handle FunctionDef, MethodDef, ClassDef, Parameter, TypeParameter,
ExceptionHandler, and PatternMatch contexts in get_type_at by
constructing a Key::Definition and looking up the type, instead of
returning None. This fixes hover on redefined functions (and other
cases where definition_at fails to find the definition).

Reviewed By: stroxler

Differential Revision: D95331714

fbshipit-source-id: 8ecfd7c52692434f346b6735c780bba292b3d810

Author: kinto0

pyrefly/lib/state/lsp.rs

@@ -912,48 +912,72 @@ impl<'a> Transaction<'a> {
                 self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context:
                     IdentifierContext::FunctionDef { docstring_range: _ }
                     | IdentifierContext::MethodDef { docstring_range: _ },
             }) => {
-                // TODO(grievejia): Handle definitions of functions
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::ClassDef { docstring_range: _ },
             }) => {
-                // TODO(grievejia): Handle definitions of classes
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::Parameter,
             }) => {
-                // TODO(grievejia): Handle definitions of params
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::TypeParameter,
             }) => {
-                // TODO(grievejia): Handle definitions of type params
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::ExceptionHandler,
             }) => {
-                // TODO(grievejia): Handle definitions of exception names
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
-                identifier: _,
+                identifier,
                 context: IdentifierContext::PatternMatch(_),
             }) => {
-                // TODO(grievejia): Handle definitions of pattern-introduced names
-                None
+                let key = Key::Definition(ShortIdentifier::new(&identifier));
+                let bindings = self.get_bindings(handle)?;
+                if !bindings.is_valid_key(&key) {
+                    return None;
+                }
+                self.get_type(handle, &key)
             }
             Some(IdentifierWithContext {
                 identifier,

pyrefly/lib/test/lsp/hover_type.rs

@@ -74,7 +74,6 @@ Hover Result: `(x: list[int], y: str, z: Literal[42]) -> list[int]`
     );
 }
 
-// TODO(kylei): redefinitions should work. they are especially common in try/except blocks
 #[test]
 fn redefinition_test() {
     let code = r#"
@@ -89,11 +88,11 @@ def f(): ...
 # main.py
 2 | def f(): ...
         ^
-Hover Result: None
+Hover Result: `() -> None`
 
 4 | def f(): ...
         ^
-Hover Result: None
+Hover Result: `() -> None`
 "#
         .trim(),
         report.trim(),