Fix numpy-style array indexing for real

[etc/vim.git] / black.py

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}: