More support for numpy tuple indexing

diff --git a/black.py b/black.py
index f305e8dd79e66d76168e16e06c3fc4302e6c61e1..b2b2db7e41192563f4eac9a37a3550584eeed437 100644 (file)
--- a/black.py
+++ b/black.py
@@ -831,7 +831,7 @@ BRACKET = {token.LPAR: token.RPAR, token.LSQB: token.RSQB, token.LBRACE: token.R
 OPENING_BRACKETS = set(BRACKET.keys())
 CLOSING_BRACKETS = set(BRACKET.values())
 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS
 OPENING_BRACKETS = set(BRACKET.keys())
 CLOSING_BRACKETS = set(BRACKET.values())
 BRACKETS = OPENING_BRACKETS | CLOSING_BRACKETS

-ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, token.COLON, STANDALONE_COMMENT}
+ALWAYS_NO_SPACE = CLOSING_BRACKETS | {token.COMMA, STANDALONE_COMMENT}

 
 
 def whitespace(leaf: Leaf) -> str:  # noqa C901
 
 
 def whitespace(leaf: Leaf) -> str:  # noqa C901


@@ -849,12 +849,18 @@ def whitespace(leaf: Leaf) -> str:  # noqa C901
         return DOUBLESPACE
 
     assert p is not None, f"INTERNAL ERROR: hand-made leaf without parent: {leaf!r}"
         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:
+        return NO
+

     prev = leaf.prev_sibling
     if not prev:
         prevp = preceding_leaf(p)
         if not prevp or prevp.type in OPENING_BRACKETS:
             return NO
 
     prev = leaf.prev_sibling
     if not prev:
         prevp = preceding_leaf(p)
         if not prevp or prevp.type in OPENING_BRACKETS:
             return NO
 

+        if t == token.COLON:
+            return SPACE if prevp.type == token.COMMA else NO
+

         if prevp.type == token.EQUAL:
             if prevp.parent and prevp.parent.type in {
                 syms.typedargslist,
         if prevp.type == token.EQUAL:
             if prevp.parent and prevp.parent.type in {
                 syms.typedargslist,


@@ -983,7 +989,7 @@ def whitespace(leaf: Leaf) -> str:  # noqa C901
 
             return NO
 
 
             return NO
 

-        elif prev.type == token.COLON:
+        else:

             return NO
 
     elif p.type == syms.atom:
             return NO
 
     elif p.type == syms.atom:

diff --git a/tests/expression.py b/tests/expression.py
index d90b7cfa799e2a0b6bc3de745c780964035ea9f3..cf37f8b98bdb1c1b63fc70497926f5794e694b20 100644 (file)
--- a/tests/expression.py
+++ b/tests/expression.py
@@ -90,6 +90,7 @@ slice[:-1]
 slice[1:]
 slice[::-1]
 numpy[:, 0:1]
 slice[1:]
 slice[::-1]
 numpy[:, 0:1]

+numpy[:, :-1]

 (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'}
 (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,6 +233,7 @@ slice[:-1]
 slice[1:]
 slice[::-1]
 numpy[:, 0:1]
 slice[1:]
 slice[::-1]
 numpy[:, 0:1]

+numpy[:, :-1]

 (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'}
 (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'}