Repository: wemake-services/wemake-python-styleguide

Issue:
WPS610 (async magic) false positive on __aiter__
### What's wrong

In Python 3.7+ a class `__aiter__` must return an async iterator object (an object that implements `__anext__`). An async generator (an async function with `yield`) is considered an async iterator object.

The following is valid and should not be flagged as an incorrect async definition.

```python
class Ex(object):
    async def __aiter__(self):
        count = 0
        while count < 50:
            count += 1
            yield count
```

The following is invalid and should be flagged as incorrect (it'll also throw a type error or somesuch if actually used in `async for`).

```python
class Ex(object):
    async def __aiter__(self):
        count = 0
        while count < 50:
            count += 1
```

### How it should be

See above. Don't throw an error.

### Flake8 version and plugins

{
  "dependencies": [],
  "platform": {
    "python_implementation": "CPython",
    "python_version": "3.8.10",
    "system": "Linux"
  },
  "plugins": [
    {
      "is_local": false,
      "plugin": "flake8-bandit",
      "version": "3.0.0"
    },
    {
      "is_local": false,
      "plugin": "flake8-broken-line",
      "version": "0.4.0"
    },
    {
      "is_local": false,
      "plugin": "flake8-bugbear",
      "version": "22.3.23"
    },
    {
      "is_local": false,
      "plugin": "flake8-class-attributes-order",
      "version": "0.1.3"
    },
    {
      "is_local": false,
      "plugin": "flake8-comprehensions",
      "version": "3.8.0"
    },
    {
      "is_local": false,
      "plugin": "flake8-darglint",
      "version": "1.8.1"
    },
    {
      "is_local": false,
      "plugin": "flake8-debugger",
      "version": "4.0.0"
    },
    {
      "is_local": false,
      "plugin": "flake8-docstrings",
      "version": "1.6.0, pydocstyle: 6.1.1"
    },
    {
      "is_local": false,
      "plugin": "flake8-eradicate",
      "version": "1.2.0"
    },
    {
      "is_local": false,
      "plugin": "flake8-string-format",
      "version": "0.3.0"
    },
    {
      "is_local": false,
      "plugin": "flake8_commas",
      "version": "2.1.0"
    },
    {
      "is_local": false,
      "plugin": "flake8_isort",
      "version": "4.1.1"
    },
    {
      "is_local": false,
      "plugin": "flake8_quotes",
      "version": "3.3.1"
    },
    {
      "is_local": false,
      "plugin": "mccabe",
      "version": "0.6.1"
    },
    {
      "is_local": false,
      "plugin": "naming",
      "version": "0.12.1"
    },
    {
      "is_local": false,
      "plugin": "pycodestyle",
      "version": "2.8.0"
    },
    {
      "is_local": false,
      "plugin": "pyflakes",
      "version": "2.4.0"
    },
    {
      "is_local": false,
      "plugin": "pylint",
      "version": "0.0.0"
    },
    {
      "is_local": false,
      "plugin": "rst-docstrings",
      "version": "0.2.5"
    },
    {
      "is_local": false,
      "plugin": "wemake_python_styleguide",
      "version": "0.16.1"
    }
  ],
  "version": "4.0.1"
}


### pip information

adblockparser==0.7
amkv==0.1.2
anyio==3.5.0
appdirs==1.4.3
apturl==0.5.2
astor==0.8.1
async-generator==1.10
async-timeout==4.0.2
asyncache==0.1.1
atomicwrites==1.1.5
attrs==19.3.0
Automat==0.8.0
bandit==1.7.4
bcrypt==3.1.7
beautifulsoup4==4.11.1
blinker==1.4
Brlapi==0.7.0
Brotli==1.0.9
cachetools==5.0.0
certifi==2019.11.28
cffi==1.14.0
chardet==4.0.0
charset-normalizer==2.0.12
Click==7.0
colorama==0.4.3
command-not-found==0.3
constantly==15.1.0
crit==0.0.1
croniter==1.3.4
cryptography==2.8
cupshelpers==1.0
cycler==0.10.0
Cython==0.29.14
darglint==1.8.1
dblatex===0.3.11py3
dbus-python==1.2.16
decorator==4.4.2
defer==1.0.6
Deprecated==1.2.13
distlib==0.3.0
distro==1.4.0
distro-info===0.23ubuntu1
dnspython==1.16.0
docutils==0.18.1
EditorConfig==0.12.2
entrypoints==0.3
eradicate==2.1.0
et-xmlfile==1.0.1
fasteners==0.14.1
ffmpeg-normalize==1.22.9
ffmpeg-progress-yield==0.2.0
filelock==3.0.12
flake8==4.0.1
flake8-bandit==3.0.0
flake8-broken-line==0.4.0
flake8-bugbear==22.3.23
flake8-class-attributes-order==0.1.3
flake8-commas==2.1.0
flake8-comprehensions==3.8.0
flake8-debugger==4.0.0
flake8-docstrings==1.6.0
flake8-eradicate==1.2.0
flake8-isort==4.1.1
flake8-polyfill==1.0.2
flake8-quotes==3.3.1
flake8-rst-docstrings==0.2.5
flake8-string-format==0.3.0
flakeheaven==0.11.1
funcparserlib==0.3.6
future==0.18.2
gevent==21.12.0
gil-load==0.4.2
gitdb==4.0.9
gitdb2==2.0.6
GitPython==3.1.27
gpg===1.13.1-unknown
graphviz==0.16
greenlet==1.1.2
gssapi==1.6.1
h11==0.12.0
h2==4.1.0
h5py==2.10.0
hpack==4.0.0
html5-parser==0.4.10
html5lib==1.0.1
httpcore==0.14.7
httplib2==0.14.0
httpx==0.22.0
hypercorn==0.13.2
hyperframe==6.0.1
hyperlink==19.0.0
idna==2.8
imageio==2.4.1
importlib-metadata==1.5.0
incremental==21.3.0
iotop==0.6
isc==2.0
isort==5.10.1
jdcal==1.0
Jinja2==2.10.1
jsbeautifier==1.10.3
Keras==2.2.4
Keras-Applications==1.0.6
Keras-Preprocessing==1.0.5
keyring==18.0.1
kiwisolver==1.0.1
language-selector==0.1
launchpadlib==1.10.13
lazr.restfulclient==0.14.2
lazr.uri==1.0.3
linecache2==1.0.0
lockfile==0.12.2
louis==3.12.0
lxml==4.5.0
macaroonbakery==1.3.1
Mako==1.1.0
Markdown==3.1.1
MarkupSafe==1.1.0
matplotlib==3.1.2
mccabe==0.6.1
memory-analyzer==0.1.2
meson==0.53.2
mock==3.0.5
monotonic==1.5
more-itertools==4.2.0
mpi4py==3.0.3
msgpack==0.6.2
multidict==6.0.2
mutagen==1.45.1
nest-asyncio==1.5.5
netaddr==0.7.19
netifaces==0.10.4
networkx==2.8
nftables==0.1
nose==1.3.7
numexpr==2.7.1
numpy==1.22.3
oauthlib==3.1.0
objgraph==3.5.0
olefile==0.46
opencv-python==4.5.5.64
openpyxl==3.0.3
outcome==1.1.0
packaging==21.3
pandas==0.25.3
parameterized==0.7.0
paramiko==2.6.0
patsy==0.5.1
pbr==5.4.5
peewee==3.14.10
pep8-naming==0.12.1
pexpect==4.6.0
Pillow==7.0.0
pluggy==0.13.0
ply==3.11
prettytable==2.1.0
priority==1.3.0
protobuf==3.6.1
psutil==5.5.1
py==1.8.1
pyasn1==0.4.2
pyasn1-modules==0.2.1
pycairo==1.16.2
pycodestyle==2.8.0
pycparser==2.19
pycrypto==2.6.1
pycryptodomex==3.14.1
pycups==1.9.73
pydocstyle==6.1.1
pydot==1.4.1
pyflakes==2.4.0
Pygments==2.11.2
PyGObject==3.36.0
pygpu==0.7.6
PyHamcrest==1.9.0
PyJWT==1.7.1
pymacaroons==0.13.0
Pympler==0.9
PyNaCl==1.3.0
pyOpenSSL==19.0.0
pyparsing==3.0.8
PyQt5==5.14.1
PyQtWebEngine==5.14.0
pyRFC3339==1.1
pytest==4.6.9
python-apt==2.0.0+ubuntu0.20.4.6
python-augeas==0.5.0
python-dateutil==2.7.3
python-debian===0.1.36ubuntu1
python-ldap==3.2.0
python-libtorrent==1.1.13
python-multipart==0.0.5
python-yubico==1.3.3
pytz==2019.3
pyusb==1.0.2
PyWavelets==1.3.0
pyxdg==0.26
PyYAML==5.3.1
pyzmq==18.1.1
qrtools==2.0
qt5reactor==0.6
QtPy==1.9.0
redis==4.2.2
reportlab==3.5.34
requests==2.27.1
requests-unixsocket==0.2.0
restructuredtext-lint==1.4.0
rfc3986==1.3.2
salt==3000.2
scikit-image==0.19.2
scipy==1.8.0
seaborn==0.10.0
SecretStorage==2.3.1
selinux==3.0
service-identity==18.1.0
simplejson==3.16.0
sip==4.19.21
six==1.14.0
slip==0.6.5
slip.dbus==0.6.5
smmap==5.0.0
smmap2==2.0.5
sniffio==1.1.0
snowballstemmer==2.0.0
sortedcontainers==2.1.0
soupsieve==1.9.5
splash==3.5
ssh-import-id==5.10
starlette==0.19.0
stevedore==1.32.0
systemd-python==234
tables==3.6.1
testfixtures==6.10.1
testresources==2.0.0
Theano==1.0.4
tifffile==2022.4.8
toml==0.10.0
tqdm==4.64.0
traceback2==1.4.0
trio==0.20.0
Twisted==22.4.0
typing-extensions==4.2.0
ubuntu-advantage-tools==27.4
ubuntu-drivers-common==0.0.0
unattended-upgrades==0.1
unittest2==1.1.0
urllib3==1.26.9
usb-creator==0.3.7
virtualenv==20.0.17
wadllib==1.3.3
wcwidth==0.1.8
webencodings==0.5.1
websockets==10.3
wemake-python-styleguide==0.16.1
Werkzeug==2.1.1
wrapt==1.14.0
wsproto==1.1.0
xdis==6.0.3
xkit==0.0.0
xlrd==1.1.0
xlwt==1.3.0
xvfbwrapper==0.2.9
yarl==1.7.2
youtube-dl==2021.12.17
yt-dlp==2022.4.8
zipp==1.0.0
zope.event==4.5.0
zope.interface==4.7.1

### OS information

Ubuntu 20.04


Discussion:
PR is welcome!

Resolving patch:
diff --git a/wemake_python_styleguide/constants.py b/wemake_python_styleguide/constants.py
--- a/wemake_python_styleguide/constants.py
+++ b/wemake_python_styleguide/constants.py
@@ -325,19 +325,18 @@
     # Allowed to be used with ``yield`` keyword:
     '__call__',
     '__iter__',
-    '__anext__',
     '__aiter__',
-    '__next__',
 })
 
 #: List of magic methods that are not allowed to be async.
 ASYNC_MAGIC_METHODS_BLACKLIST: Final = ALL_MAGIC_METHODS.difference({
     # In order of appearance on
     # https://docs.python.org/3/reference/datamodel.html#basic-customization
-    # Allowed magic methods are:
+    # Allowed async magic methods are:
     '__anext__',
     '__aenter__',
     '__aexit__',
+    '__aiter__',
     '__call__',
 })
 
diff --git a/wemake_python_styleguide/violations/oop.py b/wemake_python_styleguide/violations/oop.py
--- a/wemake_python_styleguide/violations/oop.py
+++ b/wemake_python_styleguide/violations/oop.py
@@ -445,6 +445,7 @@ class AsyncMagicMethodViolation(ASTViolation):
     Forbid certain async magic methods.
 
     We allow to make ``__anext__``, ``__aenter__``, ``__aexit__`` async.
+    We allow to make ``__aiter__`` async if it is a generator (contains yield).
     We also allow custom magic methods to be async.
 
     See
@@ -485,6 +486,7 @@ class YieldMagicMethodViolation(ASTViolation):
     Forbid ``yield`` inside of certain magic methods.
 
     We allow to make ``__iter__`` a generator.
+    We allow to make ``__aiter__`` an async generator.
     See
     :py:data:`~wemake_python_styleguide.constants.YIELD_MAGIC_METHODS_BLACKLIST`
     for the whole list of blacklisted generator magic methods.
@@ -519,7 +521,7 @@ def __init__(self):
 
     """
 
-    error_template = 'Found forbidden `yield` magic method usage'
+    error_template = 'Found forbidden `yield` magic method usage: {0}'
     code = 611
     previous_codes = {439, 435}
 
diff --git a/wemake_python_styleguide/visitors/ast/classes.py b/wemake_python_styleguide/visitors/ast/classes.py
--- a/wemake_python_styleguide/visitors/ast/classes.py
+++ b/wemake_python_styleguide/visitors/ast/classes.py
@@ -157,12 +157,22 @@ class WrongMethodVisitor(base.BaseNodeVisitor):
     _staticmethod_names: ClassVar[FrozenSet[str]] = frozenset((
         'staticmethod',
     ))
+    _special_async_iter: ClassVar[FrozenSet[str]] = frozenset((
+        '__aiter__',
+    ))
 
     def visit_any_function(self, node: types.AnyFunctionDef) -> None:
         """Checking class methods: async and regular."""
-        self._check_decorators(node)
-        self._check_bound_methods(node)
-        self._check_method_contents(node)
+        node_context = nodes.get_context(node)
+        if isinstance(node_context, ast.ClassDef):
+            self._check_decorators(node)
+            self._check_bound_methods(node)
+            self._check_yield_magic_methods(node)
+            self._check_async_magic_methods(node)
+            self._check_useless_overwritten_methods(
+                node,
+                class_name=node_context.name,
+            )
         self.generic_visit(node)
 
     def _check_decorators(self, node: types.AnyFunctionDef) -> None:
@@ -172,10 +182,6 @@ def _check_decorators(self, node: types.AnyFunctionDef) -> None:
                 self.add_violation(oop.StaticMethodViolation(node))
 
     def _check_bound_methods(self, node: types.AnyFunctionDef) -> None:
-        node_context = nodes.get_context(node)
-        if not isinstance(node_context, ast.ClassDef):
-            return
-
         if not functions.get_all_arguments(node):
             self.add_violation(
                 oop.MethodWithoutArgumentsViolation(node, text=node.name),
@@ -186,22 +192,59 @@ def _check_bound_methods(self, node: types.AnyFunctionDef) -> None:
                 oop.BadMagicMethodViolation(node, text=node.name),
             )
 
-        is_async = isinstance(node, ast.AsyncFunctionDef)
-        if is_async and access.is_magic(node.name):
-            if node.name in constants.ASYNC_MAGIC_METHODS_BLACKLIST:
+    def _check_yield_magic_methods(self, node: types.AnyFunctionDef) -> None:
+        if isinstance(node, ast.AsyncFunctionDef):
+            return
+
+        if node.name in constants.YIELD_MAGIC_METHODS_BLACKLIST:
+            if walk.is_contained(node, (ast.Yield, ast.YieldFrom)):
+                self.add_violation(
+                    oop.YieldMagicMethodViolation(node, text=node.name),
+                )
+
+    def _check_async_magic_methods(self, node: types.AnyFunctionDef) -> None:
+        if not isinstance(node, ast.AsyncFunctionDef):
+            return
+
+        if node.name in self._special_async_iter:
+            if not walk.is_contained(node, ast.Yield):  # YieldFrom not async
                 self.add_violation(
                     oop.AsyncMagicMethodViolation(node, text=node.name),
                 )
+        elif node.name in constants.ASYNC_MAGIC_METHODS_BLACKLIST:
+            self.add_violation(
+                oop.AsyncMagicMethodViolation(node, text=node.name),
+            )
 
-        self._check_useless_overwritten_methods(
-            node,
-            class_name=node_context.name,
-        )
+    def _check_useless_overwritten_methods(
+        self,
+        node: types.AnyFunctionDef,
+        class_name: str,
+    ) -> None:
+        if node.decorator_list:
+            # Any decorator can change logic and make this overwrite useful.
+            return
 
-    def _check_method_contents(self, node: types.AnyFunctionDef) -> None:
-        if node.name in constants.YIELD_MAGIC_METHODS_BLACKLIST:
-            if walk.is_contained(node, (ast.Yield, ast.YieldFrom)):
-                self.add_violation(oop.YieldMagicMethodViolation(node))
+        call_stmt = self._get_call_stmt_of_useless_method(node)
+        if call_stmt is None or not isinstance(call_stmt.func, ast.Attribute):
+            return
+
+        attribute = call_stmt.func
+        defined_method_name = node.name
+        if defined_method_name != attribute.attr:
+            return
+
+        if not super_args.is_ordinary_super_call(attribute.value, class_name):
+            return
+
+        if not function_args.is_call_matched_by_arguments(node, call_stmt):
+            return
+
+        self.add_violation(
+            oop.UselessOverwrittenMethodViolation(
+                node, text=defined_method_name,
+            ),
+        )
 
     def _get_call_stmt_of_useless_method(
         self,
@@ -236,36 +279,6 @@ def _get_call_stmt_of_useless_method(
             return stmt.value
         return None
 
-    def _check_useless_overwritten_methods(
-        self,
-        node: types.AnyFunctionDef,
-        class_name: str,
-    ) -> None:
-        if node.decorator_list:
-            # Any decorator can change logic and make this overwrite useful.
-            return
-
-        call_stmt = self._get_call_stmt_of_useless_method(node)
-        if call_stmt is None or not isinstance(call_stmt.func, ast.Attribute):
-            return
-
-        attribute = call_stmt.func
-        defined_method_name = node.name
-        if defined_method_name != attribute.attr:
-            return
-
-        if not super_args.is_ordinary_super_call(attribute.value, class_name):
-            return
-
-        if not function_args.is_call_matched_by_arguments(node, call_stmt):
-            return
-
-        self.add_violation(
-            oop.UselessOverwrittenMethodViolation(
-                node, text=defined_method_name,
-            ),
-        )
-
 
 @final
 @decorators.alias('visit_any_assign', (


Test patch:
diff --git a/tests/test_visitors/test_ast/test_classes/test_methods/test_async_yield_magic_methods.py b/tests/test_visitors/test_ast/test_classes/test_methods/test_async_yield_magic_methods.py
new file mode 100644
--- /dev/null
+++ b/tests/test_visitors/test_ast/test_classes/test_methods/test_async_yield_magic_methods.py
@@ -0,0 +1,148 @@
+import pytest
+
+from wemake_python_styleguide.violations.oop import AsyncMagicMethodViolation
+from wemake_python_styleguide.visitors.ast.classes import WrongMethodVisitor
+
+sync_method = """
+class Example(object):
+    def {0}(self, *args, **kwargs):
+        {1}
+"""
+
+async_method = """
+class Example(object):
+    async def {0}(self, *args, **kwargs):
+        {1}
+"""
+
+
+@pytest.mark.parametrize('template', [
+    sync_method,
+    async_method,
+])
+@pytest.mark.parametrize('method', [
+    '__aiter__',
+])
+@pytest.mark.parametrize('statement', [
+    'yield',
+    'yield 1',
+])
+def test_yield_is_always_allowed_in_aiter(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    template,
+    method,
+    statement,
+):
+    """Testing that the `__aiter__` can always have `yield`."""
+    tree = parse_ast_tree(template.format(method, statement))
+
+    visitor = WrongMethodVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [])
+
+
+@pytest.mark.parametrize('method', [
+    '__aiter__',
+])
+@pytest.mark.parametrize('statement', [
+    'return some_async_iterator()',
+])
+def test_wrong_async_magic_used(
+    assert_errors,
+    assert_error_text,
+    parse_ast_tree,
+    default_options,
+    method,
+    statement,
+):
+    """Testing that the method cannot be a coroutine."""
+    tree = parse_ast_tree(async_method.format(method, statement))
+
+    visitor = WrongMethodVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [AsyncMagicMethodViolation])
+    assert_error_text(visitor, method)
+
+
+@pytest.mark.parametrize('method', [
+    '__aiter__',
+])
+@pytest.mark.parametrize('statement', [
+    'yield',
+    'yield 1',
+])
+def test_correct_async_yield_magic_used(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    method,
+    statement,
+):
+    """Testing that the method can be an async generator."""
+    tree = parse_ast_tree(async_method.format(method, statement))
+
+    visitor = WrongMethodVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [])
+
+
+@pytest.mark.parametrize('method', [
+    '__aiter__',
+])
+@pytest.mark.parametrize('statement', [
+    'return some_async_iterator()',
+])
+def test_correct_sync_magic_used(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    method,
+    statement,
+):
+    """Testing that the method can be a normal method."""
+    tree = parse_ast_tree(sync_method.format(method, statement))
+
+    visitor = WrongMethodVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [])
+
+
+# Examples:
+
+correct_nested_example = """
+class Some:
+    {0}def __aiter__(self):
+        async def inner():
+            yield 1
+        return inner()
+"""
+
+
+@pytest.mark.parametrize('example', [
+    correct_nested_example,
+])
+@pytest.mark.parametrize('mode', [
+    # We don't use `mode()` fixture here, because we have a nested func.
+    '',  # sync
+    'async ',
+])
+def test_correct_examples(
+    assert_errors,
+    parse_ast_tree,
+    default_options,
+    example,
+    mode,
+):
+    """Testing specific real-life examples that should be working."""
+    tree = parse_ast_tree(example.format(mode))
+
+    visitor = WrongMethodVisitor(default_options, tree=tree)
+    visitor.run()
+
+    assert_errors(visitor, [])
diff --git a/tests/test_visitors/test_ast/test_classes/test_methods/test_yield_magic_method.py b/tests/test_visitors/test_ast/test_classes/test_methods/test_yield_magic_method.py
--- a/tests/test_visitors/test_ast/test_classes/test_methods/test_yield_magic_method.py
+++ b/tests/test_visitors/test_ast/test_classes/test_methods/test_yield_magic_method.py
@@ -27,6 +27,8 @@ def {0}(cls, *args, **kwargs):
     '__str__',
     '__aenter__',
     '__exit__',
+    '__anext__',
+    '__next__',
 ])
 @pytest.mark.parametrize('statement', [
     'yield',
@@ -35,6 +37,7 @@ def {0}(cls, *args, **kwargs):
 ])
 def test_magic_generator(
     assert_errors,
+    assert_error_text,
     parse_ast_tree,
     default_options,
     code,
@@ -48,6 +51,7 @@ def test_magic_generator(
     visitor.run()
 
     assert_errors(visitor, [YieldMagicMethodViolation])
+    assert_error_text(visitor, method)
 
 
 @pytest.mark.parametrize('code', [
@@ -90,6 +94,7 @@ def test_magic_statement(
 ])
 @pytest.mark.parametrize('method', [
     '__iter__',
+    '__aiter__',
     '__call__',
     '__custom__',
 ])
@@ -106,7 +111,7 @@ def test_iter_generator(
     method,
     statement,
 ):
-    """Testing that magic `iter` and `call` methods with `yield` are allowed."""
+    """Testing that some magic methods with `yield` are allowed."""
     tree = parse_ast_tree(code.format(method, statement))
 
     visitor = WrongMethodVisitor(default_options, tree=tree)
diff --git a/tests/test_visitors/test_ast/test_complexity/test_function/test_asserts_count.py b/tests/test_visitors/test_ast/test_complexity/test_function/test_asserts_count.py
--- a/tests/test_visitors/test_ast/test_complexity/test_function/test_asserts_count.py
+++ b/tests/test_visitors/test_ast/test_complexity/test_function/test_asserts_count.py
@@ -33,10 +33,11 @@ def test_asserts_correct_count(
     assert_errors,
     parse_ast_tree,
     code,
+    mode,
     default_options,
 ):
     """Testing that asserts counted correctly."""
-    tree = parse_ast_tree(code)
+    tree = parse_ast_tree(mode(code))
 
     visitor = FunctionComplexityVisitor(default_options, tree=tree)
     visitor.run()
@@ -54,9 +55,10 @@ def test_asserts_wrong_count(
     parse_ast_tree,
     options,
     code,
+    mode,
 ):
     """Testing that many asserts raises a warning."""
-    tree = parse_ast_tree(code)
+    tree = parse_ast_tree(mode(code))
 
     option_values = options(max_asserts=1)
     visitor = FunctionComplexityVisitor(option_values, tree=tree)