[red-knot] Emit errors for more AST nodes that are invalid (or only valid in specific contexts) in type expressions (#16822)

## Summary

Add error messages for invalid nodes in type expressions

Fixes #16816 

## Test Plan

Extend annotations/invalid.md to handle these invalid AST nodes error
messages

Author: MatthewMckee4

crates/red_knot_python_semantic/resources/mdtest/annotations/invalid.md

@@ -47,17 +47,60 @@ def _(
 ## Invalid AST nodes
 
 ```py
+def bar() -> None:
+    return None
+
 def _(
     a: 1,  # error: [invalid-type-form] "Int literals are not allowed in this context in a type expression"
     b: 2.3,  # error: [invalid-type-form] "Float literals are not allowed in type expressions"
     c: 4j,  # error: [invalid-type-form] "Complex literals are not allowed in type expressions"
     d: True,  # error: [invalid-type-form] "Boolean literals are not allowed in this context in a type expression"
-    # error: [invalid-type-form] "Bytes literals are not allowed in this context in a type expression"
-    e: int | b"foo",
+    e: int | b"foo",  # error: [invalid-type-form] "Bytes literals are not allowed in this context in a type expression"
+    f: 1 and 2,  # error: [invalid-type-form] "Boolean operations are not allowed in type expressions"
+    g: 1 or 2,  # error: [invalid-type-form] "Boolean operations are not allowed in type expressions"
+    h: (foo := 1),  # error: [invalid-type-form] "Named expressions are not allowed in type expressions"
+    i: not 1,  # error: [invalid-type-form] "Unary operations are not allowed in type expressions"
+    j: lambda: 1,  # error: [invalid-type-form] "`lambda` expressions are not allowed in type expressions"
+    k: 1 if True else 2,  # error: [invalid-type-form] "`if` expressions are not allowed in type expressions"
+    l: await 1,  # error: [invalid-type-form] "`await` expressions are not allowed in type expressions"
+    m: (yield 1),  # error: [invalid-type-form] "`yield` expressions are not allowed in type expressions"
+    n: (yield from [1]),  # error: [invalid-type-form] "`yield from` expressions are not allowed in type expressions"
+    o: 1 < 2,  # error: [invalid-type-form] "Comparison expressions are not allowed in type expressions"
+    p: bar(),  # error: [invalid-type-form] "Function calls are not allowed in type expressions"
+    q: int | f"foo",  # error: [invalid-type-form] "F-strings are not allowed in type expressions"
+    r: [1, 2, 3][1:2],  # error: [invalid-type-form] "Slices are not allowed in type expressions"
 ):
     reveal_type(a)  # revealed: Unknown
     reveal_type(b)  # revealed: Unknown
     reveal_type(c)  # revealed: Unknown
     reveal_type(d)  # revealed: Unknown
     reveal_type(e)  # revealed: int | Unknown
+    reveal_type(f)  # revealed: Unknown
+    reveal_type(g)  # revealed: Unknown
+    reveal_type(h)  # revealed: Unknown
+    reveal_type(i)  # revealed: Unknown
+    reveal_type(j)  # revealed: Unknown
+    reveal_type(k)  # revealed: Unknown
+    reveal_type(p)  # revealed: Unknown
+    reveal_type(q)  # revealed: int | Unknown
+    reveal_type(r)  # revealed: @Todo(generics)
+```
+
+## Invalid Collection based AST nodes
+
+```py
+def _(
+    a: {1: 2},  # error: [invalid-type-form] "Dict literals are not allowed in type expressions"
+    b: {1, 2},  # error: [invalid-type-form] "Set literals are not allowed in type expressions"
+    c: {k: v for k, v in [(1, 2)]},  # error: [invalid-type-form] "Dict comprehensions are not allowed in type expressions"
+    d: [k for k in [1, 2]],  # error: [invalid-type-form] "List comprehensions are not allowed in type expressions"
+    e: {k for k in [1, 2]},  # error: [invalid-type-form] "Set comprehensions are not allowed in type expressions"
+    f: (k for k in [1, 2]),  # error: [invalid-type-form] "Generator expressions are not allowed in type expressions"
+):
+    reveal_type(a)  # revealed: Unknown
+    reveal_type(b)  # revealed: Unknown
+    reveal_type(c)  # revealed: Unknown
+    reveal_type(d)  # revealed: Unknown
+    reveal_type(e)  # revealed: Unknown
+    reveal_type(f)  # revealed: Unknown
 ```

crates/red_knot_python_semantic/src/types/infer.rs

@@ -6037,16 +6037,19 @@ impl<'db> TypeInferenceBuilder<'db> {
         annotation_ty
     }
 
+    fn report_invalid_type_expression(
+        &mut self,
+        expression: &ast::Expr,
+        message: std::fmt::Arguments,
+    ) -> Type<'db> {
+        self.context
+            .report_lint(&INVALID_TYPE_FORM, expression, message);
+        Type::unknown()
+    }
+
     /// Infer the type of a type expression without storing the result.
     fn infer_type_expression_no_store(&mut self, expression: &ast::Expr) -> Type<'db> {
         // https://typing.readthedocs.io/en/latest/spec/annotations.html#grammar-token-expression-grammar-type_expression
-
-        let report_invalid_type_expression = |message: std::fmt::Arguments| {
-            self.context
-                .report_lint(&INVALID_TYPE_FORM, expression, message);
-            Type::unknown()
-        };
-
         match expression {
             ast::Expr::Name(name) => match name.ctx {
                 ast::ExprContext::Load => self
@@ -6071,47 +6074,6 @@ impl<'db> TypeInferenceBuilder<'db> {
             // https://typing.readthedocs.io/en/latest/spec/annotations.html#string-annotations
             ast::Expr::StringLiteral(string) => self.infer_string_type_expression(string),
 
-            // TODO: an Ellipsis literal *on its own* does not have any meaning in annotation
-            // expressions, but is meaningful in the context of a number of special forms.
-            ast::Expr::EllipsisLiteral(_literal) => {
-                todo_type!("ellipsis literal in type expression")
-            }
-
-            // TODO: add a subdiagnostic linking to type-expression grammar
-            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
-            ast::Expr::BytesLiteral(_) => report_invalid_type_expression(format_args!(
-                "Bytes literals are not allowed in this context in a type expression"
-            )),
-
-            // TODO: add a subdiagnostic linking to type-expression grammar
-            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
-            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
-                value: ast::Number::Int(_),
-                ..
-            }) => report_invalid_type_expression(format_args!(
-                "Int literals are not allowed in this context in a type expression"
-            )),
-
-            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
-                value: ast::Number::Float(_),
-                ..
-            }) => report_invalid_type_expression(format_args!(
-                "Float literals are not allowed in type expressions"
-            )),
-
-            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
-                value: ast::Number::Complex { .. },
-                ..
-            }) => report_invalid_type_expression(format_args!(
-                "Complex literals are not allowed in type expressions"
-            )),
-
-            // TODO: add a subdiagnostic linking to type-expression grammar
-            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
-            ast::Expr::BooleanLiteral(_) => report_invalid_type_expression(format_args!(
-                "Boolean literals are not allowed in this context in a type expression"
-            )),
-
             ast::Expr::Subscript(subscript) => {
                 let ast::ExprSubscript {
                     value,
@@ -6141,100 +6103,225 @@ impl<'db> TypeInferenceBuilder<'db> {
                 }
             }
 
-            // TODO PEP 646
-            ast::Expr::Starred(starred) => {
-                self.infer_starred_expression(starred);
-                todo_type!()
-            }
-
             // Avoid inferring the types of invalid type expressions that have been parsed from a
             // string annotation, as they are not present in the semantic index.
             _ if self.deferred_state.in_string_annotation() => Type::unknown(),
 
+            // =====================================================================================
             // Forms which are invalid in the context of annotation expressions: we infer their
             // nested expressions as normal expressions, but the type of the top-level expression is
             // always `Type::unknown` in these cases.
+            // =====================================================================================
+
+            // TODO: add a subdiagnostic linking to type-expression grammar
+            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
+            ast::Expr::BytesLiteral(_) => self.report_invalid_type_expression(
+                expression,
+                format_args!("Bytes literals are not allowed in this context in a type expression"),
+            ),
+
+            // TODO: add a subdiagnostic linking to type-expression grammar
+            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
+            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
+                value: ast::Number::Int(_),
+                ..
+            }) => self.report_invalid_type_expression(
+                expression,
+                format_args!("Int literals are not allowed in this context in a type expression"),
+            ),
+
+            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
+                value: ast::Number::Float(_),
+                ..
+            }) => self.report_invalid_type_expression(
+                expression,
+                format_args!("Float literals are not allowed in type expressions"),
+            ),
+
+            ast::Expr::NumberLiteral(ast::ExprNumberLiteral {
+                value: ast::Number::Complex { .. },
+                ..
+            }) => self.report_invalid_type_expression(
+                expression,
+                format_args!("Complex literals are not allowed in type expressions"),
+            ),
+
+            // TODO: add a subdiagnostic linking to type-expression grammar
+            // and stating that it is only valid in `typing.Literal[]` or `typing.Annotated[]`
+            ast::Expr::BooleanLiteral(_) => self.report_invalid_type_expression(
+                expression,
+                format_args!(
+                    "Boolean literals are not allowed in this context in a type expression"
+                ),
+            ),
+
             ast::Expr::BoolOp(bool_op) => {
                 self.infer_boolean_expression(bool_op);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Boolean operations are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Named(named) => {
                 self.infer_named_expression(named);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Named expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::UnaryOp(unary) => {
                 self.infer_unary_expression(unary);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Unary operations are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Lambda(lambda_expression) => {
                 self.infer_lambda_expression(lambda_expression);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("`lambda` expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::If(if_expression) => {
                 self.infer_if_expression(if_expression);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("`if` expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Dict(dict) => {
                 self.infer_dict_expression(dict);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Dict literals are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Set(set) => {
                 self.infer_set_expression(set);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Set literals are not allowed in type expressions"),
+                )
+            }
+
+            ast::Expr::DictComp(dictcomp) => {
+                self.infer_dict_comprehension_expression(dictcomp);
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Dict comprehensions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::ListComp(listcomp) => {
                 self.infer_list_comprehension_expression(listcomp);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("List comprehensions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::SetComp(setcomp) => {
                 self.infer_set_comprehension_expression(setcomp);
-                Type::unknown()
-            }
-            ast::Expr::DictComp(dictcomp) => {
-                self.infer_dict_comprehension_expression(dictcomp);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Set comprehensions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Generator(generator) => {
                 self.infer_generator_expression(generator);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Generator expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Await(await_expression) => {
                 self.infer_await_expression(await_expression);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("`await` expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Yield(yield_expression) => {
                 self.infer_yield_expression(yield_expression);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("`yield` expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::YieldFrom(yield_from) => {
                 self.infer_yield_from_expression(yield_from);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("`yield from` expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Compare(compare) => {
                 self.infer_compare_expression(compare);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Comparison expressions are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::Call(call_expr) => {
                 self.infer_call_expression(call_expr);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Function calls are not allowed in type expressions"),
+                )
             }
+
             ast::Expr::FString(fstring) => {
                 self.infer_fstring_expression(fstring);
-                Type::unknown()
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("F-strings are not allowed in type expressions"),
+                )
+            }
+
+            ast::Expr::Slice(slice) => {
+                self.infer_slice_expression(slice);
+                self.report_invalid_type_expression(
+                    expression,
+                    format_args!("Slices are not allowed in type expressions"),
+                )
+            }
+
+            // =================================================================================
+            // Branches where we probably should emit diagnostics in some context, but don't yet
+            // =================================================================================
+            ast::Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"),
+
+            ast::Expr::EllipsisLiteral(_) => {
+                todo_type!("ellipsis literal in type expression")
             }
+
             ast::Expr::List(list) => {
                 self.infer_list_expression(list);
                 Type::unknown()
             }
+
             ast::Expr::Tuple(tuple) => {
                 self.infer_tuple_expression(tuple);
                 Type::unknown()
             }
-            ast::Expr::Slice(slice) => {
-                self.infer_slice_expression(slice);
-                Type::unknown()
+
+            ast::Expr::Starred(starred) => {
+                self.infer_starred_expression(starred);
+                todo_type!("PEP 646")
             }
-            ast::Expr::IpyEscapeCommand(_) => todo!("Implement Ipy escape command support"),
         }
     }