From: Łukasz Langa <lukasz@langa.pl>
Date: Wed, 21 Mar 2018 01:17:20 +0000 (-0700)
Subject: Fix numpy-style array indexing for real
X-Git-Url: https://git.madduck.net/etc/vim.git/commitdiff_plain/5bc40707afa5fb53bbc2484ed34f69b011b98172

Fix numpy-style array indexing for real

Fixes #33
---

diff --git a/README.md b/README.md
index e4ccef2..70668f3 100644
--- a/README.md
+++ b/README.md
@@ -267,6 +267,8 @@ More details can be found in [CONTRIBUTING](CONTRIBUTING.md).
 * restored ability to format code with legacy usage of `async` as
   a name (#20, #42)
 
+* even better handling of numpy-style array indexing (#33, again)
+
 
 ### 18.3a2
 
diff --git a/black.py b/black.py
index 8a4eb05..203fbfa 100644
--- a/black.py
+++ b/black.py
@@ -872,7 +872,7 @@ def whitespace(leaf: Leaf) -> str:  # noqa C901
         return DOUBLESPACE
 
     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
-    if t == token.COLON and p.type != syms.subscript:
+    if t == token.COLON and p.type not in {syms.subscript, syms.subscriptlist}:
         return NO
 
     prev = leaf.prev_sibling
@@ -905,7 +905,7 @@ def whitespace(leaf: Leaf) -> str:  # noqa C901
                 return NO
 
         elif prevp.type == token.COLON:
-            if prevp.parent and prevp.parent.type == syms.subscript:
+            if prevp.parent and prevp.parent.type in {syms.subscript, syms.sliceop}:
                 return NO
 
         elif prevp.parent and prevp.parent.type in {syms.factor, syms.star_expr}:
diff --git a/tests/expression.py b/tests/expression.py
index cf37f8b..dbb9303 100644
--- a/tests/expression.py
+++ b/tests/expression.py
@@ -89,8 +89,25 @@ slice[:]
 slice[:-1]
 slice[1:]
 slice[::-1]
+slice[d::d + 1]
+slice[:c, c - 1]
 numpy[:, 0:1]
 numpy[:, :-1]
+numpy[0, :]
+numpy[:, i]
+numpy[0, :2]
+numpy[:N, 0]
+numpy[:2, :4]
+numpy[2:4, 1:5]
+numpy[4:, 2:]
+numpy[:, (0, 1, 2, 5)]
+numpy[0, [0]]
+numpy[:, [i]]
+numpy[1:c + 1, c]
+numpy[-(c + 1):, d]
+numpy[:, l[-2]]
+numpy[:, ::-1]
+numpy[np.newaxis, :]
 (str or None) if (sys.version_info[0] > (3,)) else (str or bytes or None)
 {'2.7': dead, '3.7': long_live or die_hard}
 {'2.7', '3.6', '3.7', '3.8', '3.9', '4.0' if gilectomy else '3.10'}
@@ -232,8 +249,25 @@ slice[:]
 slice[:-1]
 slice[1:]
 slice[::-1]
+slice[d::d + 1]
+slice[:c, c - 1]
 numpy[:, 0:1]
 numpy[:, :-1]
+numpy[0, :]
+numpy[:, i]
+numpy[0, :2]
+numpy[:N, 0]
+numpy[:2, :4]
+numpy[2:4, 1:5]
+numpy[4:, 2:]
+numpy[:, (0, 1, 2, 5)]
+numpy[0, [0]]
+numpy[:, [i]]
+numpy[1:c + 1, c]
+numpy[-(c + 1):, d]
+numpy[:, l[-2]]
+numpy[:, ::-1]
+numpy[np.newaxis, :]
 (str or None) if (sys.version_info[0] > (3,)) else (str or bytes or None)
 {'2.7': dead, '3.7': long_live or die_hard}
 {'2.7', '3.6', '3.7', '3.8', '3.9', '4.0' if gilectomy else '3.10'}