Add @parametrized handling (#27)

* Add parametrized handling

* fixup! Add parametrized handling

* Add parameters to out

* Add an example to README

Author: Tomasz Trębski

README.md

@@ -6,34 +6,40 @@
 [![Checked with mypy](http://www.mypy-lang.org/static/mypy_badge.svg)](http://mypy-lang.org/)
 [![Gitter](https://badges.gitter.im/mypy-django/Lobby.svg)](https://gitter.im/mypy-django/Lobby)
 
-
 ## Installation
 
 ```bash
 pip install pytest-mypy-plugins
 ```
 
-
 ## Usage
 
 Examples of a test case:
 
 ```yaml
 # typesafety/test_request.yml
--   case: request_object_has_user_of_type_auth_user_model
-    disable_cache: true
-    main: |
-        from django.http.request import HttpRequest
-        reveal_type(HttpRequest().user)  # N: Revealed type is 'myapp.models.MyUser'
-        # check that other fields work ok
-        reveal_type(HttpRequest().method)  # N: Revealed type is 'Union[builtins.str, None]'
-    files:
-        -   path: myapp/__init__.py
-        -   path: myapp/models.py
-            content: |
-                from django.db import models
-                class MyUser(models.Model):
-                    pass
+- case: request_object_has_user_of_type_auth_user_model
+  disable_cache: true
+  main: |
+    from django.http.request import HttpRequest
+    reveal_type(HttpRequest().user)  # N: Revealed type is 'myapp.models.MyUser'
+    # check that other fields work ok
+    reveal_type(HttpRequest().method)  # N: Revealed type is 'Union[builtins.str, None]'
+  files:
+    - path: myapp/__init__.py
+    - path: myapp/models.py
+      content: |
+        from django.db import models
+        class MyUser(models.Model):
+            pass
+- case: with_params
+  parametrized:
+    - val: 1
+      rt: builtins.int
+    - val: 1.0
+      rt: builtins.float
+  main: |
+    reveal_type({[ val }})  # N: Reveal type is '{{ rt }}'
 ```
 
 Running:

pytest_mypy_plugins/collect.py

@@ -2,13 +2,14 @@
 import platform
 import sys
 import tempfile
-from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Mapping
 
 import pytest
 import yaml
 from _pytest.config.argparsing import Parser
 from _pytest.nodes import Node
 from py._path.local import LocalPath
+import pystache
 
 from pytest_mypy_plugins import utils
 
@@ -42,6 +43,27 @@ def parse_environment_variables(env_vars: List[str]) -> Dict[str, str]:
     return parsed_vars
 
 
+def parse_parametrized(params: List[Mapping[str, Any]]) -> List[Mapping[str, Any]]:
+    if not params:
+        return [{}]
+
+    parsed_params: List[Mapping[str, Any]] = []
+    known_params = None
+    for idx, param in enumerate(params):
+        param_keys = set(sorted(param.keys()))
+        if not known_params:
+            known_params = param_keys
+        elif known_params.intersection(param_keys) != known_params:
+            raise ValueError(
+                "All parametrized entries must have same keys."
+                f'First entry is {", ".join(known_params)} but {", ".join(param_keys)} '
+                "was spotted at {idx} position",
+            )
+        parsed_params.append({k: v for k, v in param.items() if not k.startswith("__")})
+
+    return parsed_params
+
+
 class SafeLineLoader(yaml.SafeLoader):
     def construct_mapping(self, node: yaml.Node, deep: bool = False) -> None:
         mapping = super().construct_mapping(node, deep=deep)
@@ -66,37 +88,47 @@ def collect(self) -> Iterator["YamlTestItem"]:
             raise ValueError(f"Test file has to be YAML list, got {type(parsed_file)!r}.")
 
         for raw_test in parsed_file:
-            test_name = raw_test["case"]
-            if " " in test_name:
-                raise ValueError(f"Invalid test name {test_name!r}, only '[a-zA-Z0-9_]' is allowed.")
-
-            test_files = [File(path="main.py", content=raw_test["main"])]
-            test_files += parse_test_files(raw_test.get("files", []))
-
-            output_from_comments = []
-            for test_file in test_files:
-                output_lines = utils.extract_errors_from_comments(test_file.path, test_file.content.split("\n"))
-                output_from_comments.extend(output_lines)
-
-            starting_lineno = raw_test["__line__"]
-            extra_environment_variables = parse_environment_variables(raw_test.get("env", []))
-            disable_cache = raw_test.get("disable_cache", False)
-            expected_output_lines = raw_test.get("out", "").split("\n")
-            additional_mypy_config = raw_test.get("mypy_config", "")
-
-            skip = self._eval_skip(str(raw_test.get("skip", "False")))
-            if not skip:
-                yield YamlTestItem.from_parent(
-                    self,
-                    name=test_name,
-                    files=test_files,
-                    starting_lineno=starting_lineno,
-                    environment_variables=extra_environment_variables,
-                    disable_cache=disable_cache,
-                    expected_output_lines=output_from_comments + expected_output_lines,
-                    parsed_test_data=raw_test,
-                    mypy_config=additional_mypy_config,
-                )
+            test_name_prefix = raw_test["case"]
+            if " " in test_name_prefix:
+                raise ValueError(f"Invalid test name {test_name_prefix!r}, only '[a-zA-Z0-9_]' is allowed.")
+            else:
+                parametrized = parse_parametrized(raw_test.get("parametrized", []))
+
+            for params in parametrized:
+                if params:
+                    test_name_suffix = ",".join(f"{k}={v}" for k, v in params.items())
+                    test_name_suffix = f"[{test_name_suffix}]"
+                else:
+                    test_name_suffix = ""
+
+                test_name = f"{test_name_prefix}{test_name_suffix}"
+                main_file = File(path="main.py", content=pystache.render(raw_test["main"], params))
+                test_files = [main_file] + parse_test_files(raw_test.get("files", []))
+
+                output_from_comments = []
+                for test_file in test_files:
+                    output_lines = utils.extract_errors_from_comments(test_file.path, test_file.content.split("\n"))
+                    output_from_comments.extend(output_lines)
+
+                starting_lineno = raw_test["__line__"]
+                extra_environment_variables = parse_environment_variables(raw_test.get("env", []))
+                disable_cache = raw_test.get("disable_cache", False)
+                expected_output_lines = pystache.render(raw_test.get("out", ""), params).split("\n")
+                additional_mypy_config = raw_test.get("mypy_config", "")
+
+                skip = self._eval_skip(str(raw_test.get("skip", "False")))
+                if not skip:
+                    yield YamlTestItem.from_parent(
+                        self,
+                        name=test_name,
+                        files=test_files,
+                        starting_lineno=starting_lineno,
+                        environment_variables=extra_environment_variables,
+                        disable_cache=disable_cache,
+                        expected_output_lines=output_from_comments + expected_output_lines,
+                        parsed_test_data=raw_test,
+                        mypy_config=additional_mypy_config,
+                    )
 
     def _eval_skip(self, skip_if: str) -> bool:
         return eval(skip_if, {"sys": sys, "os": os, "pytest": pytest, "platform": platform})

pytest_mypy_plugins/tests/test-parametrized.yml

@@ -0,0 +1,44 @@
+---
+- case: only_main
+  parametrized:
+    - a: 1
+      revealed_type: builtins.int
+    - a: 1.0
+      revealed_type: builtins.float
+  main: |
+    a = {{ a }}
+    reveal_type(a)  # N: Revealed type is '{{ revealed_type }}'
+- case: with_extra
+  parametrized:
+    - a: 2
+      b: null
+      rt: Any
+    - a: 3
+      b: 3
+      rt: Any
+  main: |
+    import foo
+    reveal_type(foo.test({{ a }}, {{ b }}))  # N: Revealed type is '{{ rt }}'
+  files:
+    - path: foo.py
+      content: |
+        from typing import Any
+
+        def test(a: Any, b: Any) -> Any:
+          ...
+- case: with_out
+  parametrized:
+    - what: cat
+      rt: builtins.str
+    - what: dog
+      rt: builtins.str
+  main: |
+    animal = '{{ what }}'
+    reveal_type(animal)
+    try:
+      animal / 2
+    except Exception:
+      ...
+  out: |
+    main:2: note: Revealed type is '{{ rt }}'
+    main:4: error: Unsupported operand types for / ("str" and "int")

setup.py

@@ -3,7 +3,13 @@
 with open("README.md", "r") as f:
     readme = f.read()
 
-dependencies = ["pytest>=5.4.0", "mypy>=0.730", "decorator", "pyyaml"]
+dependencies = [
+    "pytest>=5.4.0",
+    "mypy>=0.730",
+    "decorator",
+    "pyyaml",
+    "pystache>=0.5.4",
+]
 
 setup(
     name="pytest-mypy-plugins",