From 30b1b1e15a0860b67a00bd1ab0e249deda05faa0 Mon Sep 17 00:00:00 2001
From: Charlie Marsh <charlie.r.marsh@gmail.com>
Date: Sun, 9 Oct 2022 17:37:19 -0400
Subject: [PATCH] Treat TypeAlias values as annotations (#377)

---
 resources/test/fixtures/F401.py | 15 +++++++++++++++
 src/ast/visitor.rs              |  7 ++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/resources/test/fixtures/F401.py b/resources/test/fixtures/F401.py
index 5d130f3755..3bb103d866 100644
--- a/resources/test/fixtures/F401.py
+++ b/resources/test/fixtures/F401.py
@@ -67,7 +67,22 @@ c = cast("Vegetable", b)
 Field = lambda default=MISSING: field(default=default)
 
 
+# Test: access a sub-importation via an alias.
 import pyarrow as pa
 import pyarrow.csv
 
 print(pa.csv.read_csv("test.csv"))
+
+
+# Test: referencing an import via TypeAlias.
+import sys
+
+import numpy as np
+
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+
+CustomInt: TypeAlias = "np.int8 | np.int16"
diff --git a/src/ast/visitor.rs b/src/ast/visitor.rs
index d47fbdaaa2..911e55a43a 100644
--- a/src/ast/visitor.rs
+++ b/src/ast/visitor.rs
@@ -1,3 +1,4 @@
+use crate::ast::helpers::match_name_or_attr;
 use rustpython_parser::ast::{
     Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Constant, Excepthandler,
     ExcepthandlerKind, Expr, ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern,
@@ -148,7 +149,11 @@ pub fn walk_stmt<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, stmt: &'a Stmt) {
         } => {
             visitor.visit_annotation(annotation);
             if let Some(expr) = value {
-                visitor.visit_expr(expr);
+                if match_name_or_attr(annotation, "TypeAlias") {
+                    visitor.visit_annotation(expr);
+                } else {
+                    visitor.visit_expr(expr);
+                }
             }
             visitor.visit_expr(target);
         }