check cancel status of tokens before all when doing sum

pomponchik · Aug 8, 2024 · ab318ae · ab318ae
1 parent 506a55c
commit ab318ae
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/cantok/tokens/abstract/abstract_token.py b/cantok/tokens/abstract/abstract_token.py
@@ -55,6 +55,9 @@ def __add__(self, item: 'AbstractToken') -> 'AbstractToken':
 
         from cantok import SimpleToken, DefaultToken, TimeoutToken
 
+        if self._cancelled or item._cancelled:
+            return SimpleToken(cancelled=True)
+
         nested_tokens = []
         container_token: Optional[AbstractToken] = None
 
@@ -79,9 +82,7 @@ def __add__(self, item: 'AbstractToken') -> 'AbstractToken':
                         return self
 
         for token in self, item:
-            if token._cancelled:
-                return SimpleToken(cancelled=True)
-            elif isinstance(token, SimpleToken) and getrefcount(token) < 6:
+            if isinstance(token, SimpleToken) and getrefcount(token) < 6:
                 nested_tokens.extend(token.tokens)
             elif isinstance(token, DefaultToken):
                 pass