From d29f81f89ffa6f21f15bafa905f80c44ce2f32b7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=20=D0=9B=D0=B5=D0=B1=D0=B5=D0=B4?=
 =?UTF-8?q?=D0=B5=D0=B2?= <i.lebedev@bestdoctor.ru>
Date: Sun, 12 Jan 2020 11:38:49 +0300
Subject: [PATCH] FIX #7: add async with support and with variables nodes
 validation

---
 flake8_expression_complexity/__init__.py  |  2 +-
 flake8_expression_complexity/utils/ast.py | 11 +++++++----
 requirements_dev.txt                      |  2 +-
 setup.cfg                                 |  2 +-
 tests/test_complexity.py                  |  2 +-
 tests/test_files/long_expressions.py      |  7 ++++++-
 6 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/flake8_expression_complexity/__init__.py b/flake8_expression_complexity/__init__.py
index eead319..fa9c4ec 100644
--- a/flake8_expression_complexity/__init__.py
+++ b/flake8_expression_complexity/__init__.py
@@ -1 +1 @@
-__version__ = '0.0.5'
+__version__ = '0.0.6'
diff --git a/flake8_expression_complexity/utils/ast.py b/flake8_expression_complexity/utils/ast.py
index d0e8afb..148676c 100644
--- a/flake8_expression_complexity/utils/ast.py
+++ b/flake8_expression_complexity/utils/ast.py
@@ -5,17 +5,20 @@
 
 def iterate_over_expressions(node: ast.AST) -> Iterable[ast.AST]:
     additionals_subnodes_info: List[Tuple[Tuple, Callable]] = [
-        ((ast.If, ast.While), lambda n: n.test),
-        ((ast.For, ), lambda n: n.iter),
+        ((ast.If, ast.While), lambda n: [n.test]),
+        ((ast.For, ), lambda n: [n.iter]),
+        ((ast.With, ast.AsyncWith), lambda n: [s.context_expr for s in n.items]),
     ]
     nodes_with_subnodes = (
         ast.FunctionDef, ast.AsyncFunctionDef,
         ast.If, ast.For, ast.Module,
-        ast.ClassDef, ast.Try, ast.With, ast.While,
+        ast.ClassDef, ast.Try, ast.With, ast.AsyncWith,
+        ast.While,
     )
     for bases, subnodes_getter in additionals_subnodes_info:
         if isinstance(node, bases):
-            yield subnodes_getter(node)
+            for subitem in subnodes_getter(node):
+                yield subitem
     nodes_to_iter = (
         _get_try_node_children(node)
         if isinstance(node, ast.Try)
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 9b7dd1f..c1a070e 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -2,7 +2,7 @@ mypy==0.740
 mypy-extensions==0.4.1
 typing==3.6.4
 typing-extensions==3.7.4
-flake8==3.7.8
+flake8==3.7.9
 flake8-annotations-complexity==0.0.2
 flake8-blind-except==0.1.1
 flake8-broken-line==0.1.1
diff --git a/setup.cfg b/setup.cfg
index 86610cb..c2215d7 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -3,7 +3,7 @@ ignore = D, W503
 max-line-length = 100
 use_class_attributes_order_strict_mode = True
 max_function_length = 50
-max-complexity = 6
+max-complexity = 7
 
 
 [mypy]
diff --git a/tests/test_complexity.py b/tests/test_complexity.py
index b576802..da8cf3d 100644
--- a/tests/test_complexity.py
+++ b/tests/test_complexity.py
@@ -3,4 +3,4 @@
 
 def test_fails():
     errors = run_validator_for_test_file('long_expressions.py', max_expression_compexity=3)
-    assert len(errors) == 3
+    assert len(errors) == 4
diff --git a/tests/test_files/long_expressions.py b/tests/test_files/long_expressions.py
index 102696a..15eb351 100644
--- a/tests/test_files/long_expressions.py
+++ b/tests/test_files/long_expressions.py
@@ -41,8 +41,13 @@ async def foo():
 
 
 async def bar():
-    return 'bar'
+    async with foo:
+        return 'bar'
 
 
 weird_container = []
 sublist = weird_container[10:datetime.datetime.today(), None]
+
+
+with weird_container[10:str(datetime.datetime.today().date())[:100], None]:
+    pass