Fix numpy-style array indexing for real

Fixes #33

diff --git a/README.md b/README.md
index e4ccef20fc75df350b9a4c3f9fddb3ba8b4a4793..70668f3180a0d7d18d2074dba9dcb81a1d259006 100644 (file)
--- 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 8a4eb05e5ad71cf04976c913a0f8bac6631bae7f..203fbfa9055be595a358e94bd1c16b9a3911f23f 100644 (file)
--- 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 cf37f8b98bdb1c1b63fc70497926f5794e694b20..dbb9303673eb2b98feb69ccdec018ecc28c5b09e 100644 (file)
--- 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'}