From b2882fd6be4dded74a2f1be36c2e4c2d18e3348d Mon Sep 17 00:00:00 2001
From: Eric Traut <erictr@microsoft.com>
Date: Fri, 2 Oct 2020 10:55:24 -0700
Subject: [PATCH] Added support for the `__all__ += <module>.__all__` idiom for
 mutating the `__all__` value. This idiom is used by numpy.

---
 docs/typed-libraries.md                         |  1 +
 .../pyright-internal/src/analyzer/binder.ts     | 17 ++++++++++++++++-
 .../import.pytyped.dunderAll.fourslash.ts       | 10 ++++++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/docs/typed-libraries.md b/docs/typed-libraries.md
index 50c26f5a9..60dd9889e 100644
--- a/docs/typed-libraries.md
+++ b/docs/typed-libraries.md
@@ -40,6 +40,7 @@ The following idioms are supported for defining the values contained within `__a
 
 * `__all__ = ['a', b']`
 * `__all__ += ['a', b']`
+* `__all__ += submodule.__all__`
 * `__all__.extend(['a', b'])`
 * `__all__.append('a')`
 * `__all__.remove('a')`
diff --git a/packages/pyright-internal/src/analyzer/binder.ts b/packages/pyright-internal/src/analyzer/binder.ts
index c4b143782..f15819cb9 100644
--- a/packages/pyright-internal/src/analyzer/binder.ts
+++ b/packages/pyright-internal/src/analyzer/binder.ts
@@ -734,8 +734,10 @@ export class Binder extends ParseTreeWalker {
         this._bindPossibleTupleNamedTarget(node.destExpression);
         this._createAssignmentTargetFlowNodes(node.destExpression, /* walkTargets */ false, /* unbound */ false);
 
-        // Is this an assignment to dunder all?
+        // Is this an assignment to dunder all of the form
+        // __all__ += <expression>?
         if (
+            node.operator === OperatorType.AddEqual &&
             this._currentScope.type === ScopeType.Module &&
             node.leftExpression.nodeType === ParseNodeType.Name &&
             node.leftExpression.value === '__all__'
@@ -743,6 +745,7 @@ export class Binder extends ParseTreeWalker {
             const expr = node.rightExpression;
 
             if (expr.nodeType === ParseNodeType.List) {
+                // Is this the form __all__ += ["a", "b"]?
                 expr.entries.forEach((listEntryNode) => {
                     if (
                         listEntryNode.nodeType === ParseNodeType.StringList &&
@@ -752,6 +755,18 @@ export class Binder extends ParseTreeWalker {
                         this._dunderAllNames?.push(listEntryNode.strings[0].value);
                     }
                 });
+            } else if (
+                expr.nodeType === ParseNodeType.MemberAccess &&
+                expr.leftExpression.nodeType === ParseNodeType.Name &&
+                expr.memberName.value === '__all__'
+            ) {
+                // Is this using the form "__all__ += <mod>.__all__"?
+                const namesToAdd = this._getDunderAllNamesFromImport(expr.leftExpression.value);
+                if (namesToAdd) {
+                    namesToAdd.forEach((name) => {
+                        this._dunderAllNames?.push(name);
+                    });
+                }
             }
         }
 
diff --git a/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts
index 2bd30cbc2..445936313 100644
--- a/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts
+++ b/packages/pyright-internal/src/tests/fourslash/import.pytyped.dunderAll.fourslash.ts
@@ -7,6 +7,7 @@
 // @filename: testpkg/__init__.py
 // @library: true
 //// from . import submod
+//// from .submod2 import *
 //// from submod import foofoofoo5, foofoofoo6, foofoofoo7, foofoofoo8
 //// foofoofoo0: int = 0
 //// foofoofoo1: int = 1
@@ -20,6 +21,7 @@
 //// __all__.remove("foofoofoo1")
 //// __all__.remove("foofoofoo6")
 //// __all__.append("foofoofoo0")
+//// __all__ += submod2.__all__
 
 // @filename: testpkg/submod.py
 // @library: true
@@ -31,6 +33,11 @@
 //// __all__ += ["foofoofoo6"]
 //// __all__.extend(["foofoofoo7"])
 
+// @filename: testpkg/submod2.py
+// @library: true
+//// foofoofoo9: int = 9
+//// __all__ = ["foofoofoo9"]
+
 // @filename: .src/test.py
 //// from testpkg import *
 //// foofoofoo[|/*marker1*/|]
@@ -56,6 +63,9 @@ await helper.verifyCompletion('exact', {
             {
                 label: 'foofoofoo7',
             },
+            {
+                label: 'foofoofoo9',
+            },
         ],
     },
 });