From da80481a448746648433556863f9024c62778bab Mon Sep 17 00:00:00 2001
From: Ofek Inger <ofeki@neosec.com>
Date: Fri, 1 Nov 2024 10:37:44 -0400
Subject: [PATCH] Fixed type annotations in assignments not creating symbols

---
 .../src/TypeStubExtendedWriter.ts             | 57 +++++++++++++++++++
 packages/pyright-scip/src/treeVisitor.ts      | 28 ++++++++-
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/packages/pyright-scip/src/TypeStubExtendedWriter.ts b/packages/pyright-scip/src/TypeStubExtendedWriter.ts
index 68423cd0f..7c4bc7d3c 100644
--- a/packages/pyright-scip/src/TypeStubExtendedWriter.ts
+++ b/packages/pyright-scip/src/TypeStubExtendedWriter.ts
@@ -19,6 +19,7 @@ import {
     ParameterCategory,
     ParameterNode,
     ParseNodeType,
+    TypeAnnotationNode,
 } from 'pyright-internal/parser/parseNodes';
 import * as TypeUtils from 'pyright-internal/analyzer/typeUtils';
 import * as ParseTreeUtils from 'pyright-internal/analyzer/parseTreeUtils';
@@ -143,6 +144,62 @@ export class TypeStubExtendedWriter extends TypeStubWriter {
         return true;
     }
 
+    override visitTypeAnnotation(node: TypeAnnotationNode): boolean {
+        let isTypeAlias = false;
+        let line = '';
+
+        // Handle cases where left expression is a simple name assignment (e.g., "a = 1").
+        if (node.valueExpression.nodeType === ParseNodeType.Name && node.parent?.nodeType === ParseNodeType.Assignment) {
+            // TODO: Handle "__all__" as a special case.
+            // if (leftExpr.value === '__all__') {
+            //     if (this._functionNestCount === 0 && this._ifNestCount === 0) {
+            //         this._emittedSuite = true;
+            //
+            //         line = this._printExpression(leftExpr);
+            //         line += ' = ';
+            //         line += this._printExpression(node.rightExpression);
+            //         this._emitLine(line);
+            //     }
+            //
+            //     return false;
+            // }
+
+            const valueType = this.evaluator.getType(node.valueExpression);
+
+            if (node.parent.typeAnnotationComment) {
+                line += this._printExpression(node.parent.typeAnnotationComment, /* treatStringsAsSymbols */ true);
+            } else if (valueType) {
+                line += TypeUtils.getFullNameOfType(valueType);
+            }
+
+            if (valueType?.typeAliasInfo) {
+                isTypeAlias = true;
+            } else if (node.parent.rightExpression.nodeType === ParseNodeType.Call) {
+                // Special-case TypeVar, TypeVarTuple, ParamSpec and NewType calls.
+                // Treat them like type aliases.
+                const callBaseType = this.evaluator.getType(node.parent.rightExpression.leftExpression);
+                if (
+                    callBaseType &&
+                    isInstantiableClass(callBaseType) &&
+                    ClassType.isBuiltIn(callBaseType, ['TypeVar', 'TypeVarTuple', 'ParamSpec', 'NewType'])
+                ) {
+                    isTypeAlias = true;
+                }
+            }
+
+            if (line && isTypeAlias) {
+                line += ' = ';
+                line += this._printExpression(node.parent.rightExpression);
+            }
+        }
+
+        if (line) {
+            this.docstrings.set(node.id, [line]);
+        }
+
+        return true;
+    }
+
     override visitAssignment(node: AssignmentNode): boolean {
         let isTypeAlias = false;
         let line = '';
diff --git a/packages/pyright-scip/src/treeVisitor.ts b/packages/pyright-scip/src/treeVisitor.ts
index 69115d90e..306efa59a 100644
--- a/packages/pyright-scip/src/treeVisitor.ts
+++ b/packages/pyright-scip/src/treeVisitor.ts
@@ -279,6 +279,24 @@ export class TreeVisitor extends ParseTreeWalker {
                     documentation: _formatHover(hoverResult!),
                 })
             );
+        } 
+        else if (node.parent?.nodeType == ParseNodeType.Assignment && node.valueExpression.nodeType === ParseNodeType.Name) {
+            this._docstringWriter.visitTypeAnnotation(node);
+
+            let documentation = [];
+
+            let assignmentDoc = this._docstringWriter.docstrings.get(node.id);
+            if (assignmentDoc) {
+                documentation.push('```python\n' + assignmentDoc.join('\n') + '\n```');
+            }
+
+            // node.typeAnnotationComment
+            this.document.symbols.push(
+                new scip.SymbolInformation({
+                    symbol: this.getScipSymbol(node).value,
+                    documentation,
+                })
+            );
         }
 
         return true;
@@ -1164,8 +1182,14 @@ export class TreeVisitor extends ParseTreeWalker {
                         return ScipSymbol.local(this.counter.next());
                     }
                 }
-
-                return Symbols.makeTerm(this.getScipSymbol(enclosingSuite || parent), (node as NameNode).value);
+                
+                let symbol = null;
+                if (parent.nodeType === ParseNodeType.TypeAnnotation)
+                    symbol = this.getScipSymbol(enclosingSuite || parent);
+                else
+                    symbol = Symbols.makeTerm(this.getScipSymbol(enclosingSuite || parent), (node as NameNode).value);
+                
+                return symbol;
             }
             case ParseNodeType.TypeAnnotation: {
                 switch (node.valueExpression.nodeType) {