├── .github └── dependabot.yml ├── .gitignore ├── .travis.yml ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── appveyor.yml ├── example ├── .gitignore ├── README.md ├── driver.py ├── example_conftest.py ├── gcd.py └── test_gcd.py ├── mypy.ini ├── pyannotate_runtime ├── __init__.py ├── collect_types.py └── tests │ ├── __init__.py │ └── test_collect_types.py ├── pyannotate_tools ├── __init__.py ├── annotations │ ├── __init__.py │ ├── __main__.py │ ├── infer.py │ ├── main.py │ ├── parse.py │ ├── tests │ │ ├── __init__.py │ │ ├── dundermain_test.py │ │ ├── infer_test.py │ │ ├── main_test.py │ │ ├── parse_test.py │ │ └── types_test.py │ └── types.py └── fixes │ ├── __init__.py │ ├── fix_annotate.py │ ├── fix_annotate_json.py │ └── tests │ ├── __init__.py │ ├── test_annotate_json_py2.py │ ├── test_annotate_json_py3.py │ ├── test_annotate_py2.py │ └── test_annotate_py3.py ├── requirements.txt ├── setup.cfg ├── setup.py └── tests └── integration_test.py /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: pip 4 | directory: "/" 5 | schedule: 6 | interval: daily 7 | open-pull-requests-limit: 10 8 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /.cache 2 | /*.egg-info 3 | /build 4 | /dist 5 | *.pyc 6 | __pycache__ 7 | /env* 8 | *~ 9 | .mypy_cache/ 10 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | cache: pip 3 | matrix: 4 | include: 5 | - python: '2.7' 6 | - python: '3.4' 7 | - python: '3.5.1' # Special, it doesn't have typing.Text 8 | dist: trusty # needed because xenial doesn't have 3.5.1 9 | - python: '3.5' # Latest, e.g. 3.5.4 10 | - python: '3.6' 11 | env: USE_MYPY=true 12 | - python: '3.7' 13 | dist: xenial # needed because Python 3.7 is broken on travis CI Trusty 14 | sudo: true 15 | env: USE_MYPY=true 16 | - python: '3.8-dev' 17 | dist: xenial # needed because Python 3.8 is broken on travis CI Trusty 18 | sudo: true 19 | env: USE_MYPY=true 20 | allow_failures: 21 | - python: 3.8-dev 22 | 23 | install: 24 | - pip install -r requirements.txt 25 | - if [[ $USE_MYPY == true ]]; then pip install -U mypy; fi 26 | script: 27 | - pytest 28 | - if [[ $USE_MYPY == true ]]; then mypy pyannotate_*; fi 29 | -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- 1 | First-time contributors: Please fill out the Dropbox Contributor 2 | License Agreement (CLA) at https://opensource.dropbox.com/cla/ 3 | (we currently check this manually, so we apologize for delays). 4 | 5 | Everyone: 6 | 7 | - Please run the tests (`pytest`) and make sure they pass. 8 | - Please add tests for the bug/feature you are fixing/adding. 9 | - Please follow PEP 8 for coding style. 10 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright (c) 2017 Dropbox, Inc. 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | graft example 2 | graft pyannotate_tools 3 | graft pyannotate_runtime 4 | include README.md 5 | include CONTRIBUTING.md 6 | include LICENSE 7 | include setup.py 8 | include setup.cfg 9 | include mypy.ini 10 | include requirements.txt 11 | global-exclude @* 12 | global-exclude *~ 13 | global-exclude *.pyc 14 | global-exclude *.json 15 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | PyAnnotate: Auto-generate PEP-484 annotations 2 | ============================================= 3 | 4 | Insert annotations into your source code based on call arguments and 5 | return types observed at runtime. 6 | 7 | For license and copyright see the end of this file. 8 | 9 | Blog post: http://mypy-lang.blogspot.com/2017/11/dropbox-releases-pyannotate-auto.html 10 | 11 | How to use 12 | ========== 13 | 14 | See also the example directory. 15 | 16 | Phase 1: Collecting types at runtime 17 | ------------------------------------ 18 | 19 | - Install the usual way (see "red tape" section below) 20 | - Add `from pyannotate_runtime import collect_types` to your test 21 | - Early in your test setup, call `collect_types.init_types_collection()` 22 | - Bracket your test execution between calls to `collect_types.start()` and 23 | `collect_types.stop()` (or use the context manager below) 24 | - When done, call `collect_types.dump_stats(filename)` 25 | 26 | All calls between the `start()` and `stop()` calls will be analyzed 27 | and the observed types will be written (in JSON form) to the filename 28 | you pass to `dump_stats()`. You can have multiple start/stop pairs 29 | per dump call. 30 | 31 | If you'd like to automatically collect types when you run `pytest`, 32 | see `example/example_conftest.py` and `example/README.md`. 33 | 34 | Instead of using `start()` and `stop()` you can also use a context 35 | manager: 36 | ``` 37 | collect_types.init_types_collection() 38 | with collect_types.collect(): 39 | 40 | collect_types.dump_stats() 41 | ``` 42 | 43 | Phase 2: Inserting types into your source code 44 | ---------------------------------------------- 45 | 46 | The command-line tool `pyannotate` can add annotations into your 47 | source code based on the annotations collected in phase 1. The key 48 | arguments are: 49 | 50 | - Use `--type-info FILE` to tell it the file you passed to `dump_stats()` 51 | - Positional arguments are source files you want to annotate 52 | - With no other flags the tool will print a diff indicating what it 53 | proposes to do but won't do anything. Review the output. 54 | - Add `-w` to make the tool actually update your files. 55 | (Use git or some other way to keep a backup.) 56 | 57 | At this point you should probably run mypy and iterate. You probably 58 | will have to tweak the changes to make mypy completely happy. 59 | 60 | Notes and tips 61 | -------------- 62 | 63 | - It's best to do one file at a time, at least until you're 64 | comfortable with the tool. 65 | - The tool doesn't touch functions that already have an annotation. 66 | - The tool can generate either of: 67 | - type comments, i.e. Python 2 style annotations 68 | - inline type annotations, i.e. Python 3 style annotations, using `--py3` in v1.0.7+ 69 | 70 | Red tape 71 | ======== 72 | 73 | Installation 74 | ------------ 75 | 76 | This should work for Python 2.7 as well as for Python 3.4 and higher. 77 | 78 | ``` 79 | pip install pyannotate 80 | ``` 81 | 82 | This installs several items: 83 | 84 | - A runtime module, pyannotate_runtime/collect_types.py, which collects 85 | and dumps types observed at runtime using a profiling hook. 86 | 87 | - A library package, pyannotate_tools, containing code that can read the 88 | data dumped by the runtime module and insert annotations into your 89 | source code. 90 | 91 | - An entry point, pyannotate, which runs the library package on your files. 92 | 93 | For dependencies, see setup.py and requirements.txt. 94 | 95 | Testing etc. 96 | ------------ 97 | 98 | To run the unit tests, use pytest: 99 | 100 | ``` 101 | pytest 102 | ``` 103 | 104 | TO DO 105 | ----- 106 | 107 | We'd love your help with some of these issues: 108 | 109 | - Better documentation. 110 | - Python 3 code generation. 111 | - Refactor the tool modules (currently its legacy architecture shines through). 112 | 113 | Acknowledgments 114 | --------------- 115 | 116 | The following people contributed significantly to this tool: 117 | 118 | - Tony Grue 119 | - Sergei Vorobev 120 | - Jukka Lehtosalo 121 | - Guido van Rossum 122 | 123 | Licence etc. 124 | ------------ 125 | 126 | 1. License: Apache 2.0. 127 | 2. Copyright attribution: Copyright (c) 2017 Dropbox, Inc. 128 | 3. External contributions to the project should be subject to 129 | Dropbox's Contributor License Agreement (CLA): 130 | https://opensource.dropbox.com/cla/ 131 | -------------------------------------------------------------------------------- /appveyor.yml: -------------------------------------------------------------------------------- 1 | environment: 2 | matrix: 3 | - PYTHON: "C:\\Python27" 4 | - PYTHON: "C:\\Python36-x64" 5 | 6 | build: off 7 | 8 | install: 9 | - "%PYTHON%\\python.exe -m pip install -r requirements.txt" 10 | - "%PYTHON%\\python.exe -m pip list" 11 | 12 | test_script: 13 | - "%PYTHON%\\python.exe -m pytest" 14 | -------------------------------------------------------------------------------- /example/.gitignore: -------------------------------------------------------------------------------- 1 | type_info.json 2 | -------------------------------------------------------------------------------- /example/README.md: -------------------------------------------------------------------------------- 1 | PyAnnotate example 2 | ================== 3 | 4 | To play with this example, first install PyAnnotate: 5 | 6 | ``` 7 | pip install pyannotate 8 | ``` 9 | 10 | Then run the driver.py file: 11 | 12 | ``` 13 | python driver.py 14 | ``` 15 | 16 | Expected contents of type_info.json (after running driver.py): 17 | 18 | ``` 19 | [ 20 | { 21 | "path": "gcd.py", 22 | "line": 1, 23 | "func_name": "main", 24 | "type_comments": [ 25 | "() -> None" 26 | ], 27 | "samples": 1 28 | }, 29 | { 30 | "path": "gcd.py", 31 | "line": 5, 32 | "func_name": "gcd", 33 | "type_comments": [ 34 | "(int, int) -> int" 35 | ], 36 | "samples": 2 37 | } 38 | ] 39 | ``` 40 | 41 | Now run the pyannotate tool, like this (note the -w flag -- without 42 | this it won't update the file): 43 | 44 | ``` 45 | pyannotate -w gcd.py 46 | ``` 47 | 48 | Expected output: 49 | 50 | ``` 51 | Refactored gcd.py 52 | --- gcd.py (original) 53 | +++ gcd.py (refactored) 54 | @@ -1,8 +1,10 @@ 55 | def main(): 56 | + # type: () -> None 57 | print(gcd(15, 10)) 58 | print(gcd(45, 12)) 59 | 60 | def gcd(a, b): 61 | + # type: (int, int) -> int 62 | while b: 63 | a, b = b, a%b 64 | return a 65 | Files that were modified: 66 | gcd.py 67 | ``` 68 | 69 | Alternative, using pytest 70 | ------------------------- 71 | 72 | For pytest users, the example_conftest.py file shows how to 73 | automatically configures pytest to collect types when running tests. 74 | The test_gcd.py file contains a simple test to demonstrate this. Copy 75 | the contents of example_conftest.py to your conftest.py file and run 76 | pytest; it will then generate a type_info.json file like the one 77 | above. 78 | -------------------------------------------------------------------------------- /example/driver.py: -------------------------------------------------------------------------------- 1 | from gcd import main 2 | from pyannotate_runtime import collect_types 3 | 4 | if __name__ == '__main__': 5 | collect_types.init_types_collection() 6 | with collect_types.collect(): 7 | main() 8 | collect_types.dump_stats('type_info.json') 9 | -------------------------------------------------------------------------------- /example/example_conftest.py: -------------------------------------------------------------------------------- 1 | # Configuration for pytest to automatically collect types. 2 | # Thanks to Guilherme Salgado. 3 | 4 | import pytest 5 | 6 | 7 | def pytest_collection_finish(session): 8 | """Handle the pytest collection finish hook: configure pyannotate. 9 | 10 | Explicitly delay importing `collect_types` until all tests have 11 | been collected. This gives gevent a chance to monkey patch the 12 | world before importing pyannotate. 13 | """ 14 | from pyannotate_runtime import collect_types 15 | collect_types.init_types_collection() 16 | 17 | 18 | @pytest.fixture(autouse=True) 19 | def collect_types_fixture(): 20 | from pyannotate_runtime import collect_types 21 | collect_types.start() 22 | yield 23 | collect_types.stop() 24 | 25 | 26 | def pytest_sessionfinish(session, exitstatus): 27 | from pyannotate_runtime import collect_types 28 | collect_types.dump_stats("type_info.json") 29 | -------------------------------------------------------------------------------- /example/gcd.py: -------------------------------------------------------------------------------- 1 | def main(): 2 | print(gcd(15, 10)) 3 | print(gcd(45, 12)) 4 | 5 | def gcd(a, b): 6 | while b: 7 | a, b = b, a%b 8 | return a 9 | -------------------------------------------------------------------------------- /example/test_gcd.py: -------------------------------------------------------------------------------- 1 | # Tests for gcd function. 2 | 3 | from gcd import gcd 4 | 5 | def test_gcd(): 6 | assert gcd(5, 10) == 5 7 | assert gcd(12, 45) == 3 8 | -------------------------------------------------------------------------------- /mypy.ini: -------------------------------------------------------------------------------- 1 | [mypy] 2 | ignore_missing_imports = True 3 | strict_optional = True 4 | -------------------------------------------------------------------------------- /pyannotate_runtime/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dropbox/pyannotate/a7a46f394f0ba91a1b5fbf657e2393af542969ae/pyannotate_runtime/__init__.py -------------------------------------------------------------------------------- /pyannotate_runtime/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dropbox/pyannotate/a7a46f394f0ba91a1b5fbf657e2393af542969ae/pyannotate_runtime/tests/__init__.py -------------------------------------------------------------------------------- /pyannotate_runtime/tests/test_collect_types.py: -------------------------------------------------------------------------------- 1 | """Tests for collect_types""" 2 | from __future__ import ( 3 | absolute_import, 4 | division, 5 | print_function, 6 | ) 7 | 8 | import contextlib 9 | import json 10 | import os 11 | import sched 12 | import sys 13 | import time 14 | import unittest 15 | from collections import namedtuple 16 | from threading import Thread 17 | 18 | from six import PY2 19 | from typing import ( 20 | Any, 21 | Dict, 22 | Iterator, 23 | List, 24 | Optional, 25 | Tuple, 26 | Union, 27 | ) 28 | try: 29 | from typing import Text 30 | except ImportError: 31 | # In Python 3.5.1 stdlib, typing.py does not define Text 32 | Text = str # type: ignore 33 | 34 | from pyannotate_runtime import collect_types 35 | 36 | # A bunch of random functions and classes to test out type collection 37 | # Disable a whole bunch of lint warnings for simplicity 38 | 39 | # pylint:disable=invalid-name 40 | # pylint:disable=blacklisted-name 41 | # pylint:disable=missing-docstring 42 | 43 | FooNamedTuple = namedtuple('FooNamedTuple', 'foo bar') 44 | 45 | 46 | def print_int(i): 47 | # type: (Any) -> Any 48 | print(i) 49 | 50 | 51 | def noop_dec(a): 52 | # type: (Any) -> Any 53 | return a 54 | 55 | def discard(a): 56 | # type: (Any) -> None 57 | pass 58 | 59 | @noop_dec 60 | class FoosParent(object): 61 | pass 62 | 63 | 64 | class FooObject(FoosParent): 65 | class FooNested(object): 66 | pass 67 | 68 | 69 | class FooReturn(FoosParent): 70 | pass 71 | 72 | 73 | class WorkerClass(object): 74 | 75 | def __init__(self, special_num, foo): 76 | # type: (Any, Any) -> None 77 | self._special_num = special_num 78 | self._foo = foo 79 | 80 | @noop_dec 81 | def do_work(self, i, haz): 82 | # type: (Any, Any) -> Any 83 | print_int(i) 84 | return EOFError() 85 | 86 | @classmethod 87 | def do_work_clsmthd(cls, i, haz=None): 88 | # type: (Any, Any) -> Any 89 | print_int(i) 90 | return EOFError() 91 | 92 | 93 | class EventfulHappenings(object): 94 | 95 | def __init__(self): 96 | # type: () -> None 97 | self.handlers = [] # type: Any 98 | 99 | def add_handler(self, handler): 100 | # type: (Any) -> Any 101 | self.handlers.append(handler) 102 | 103 | def something_happened(self, a, b): 104 | # type: (Any, Any) -> Any 105 | for h in self.handlers: 106 | h(a, b) 107 | return 1999 108 | 109 | 110 | # A class that is old style under python 2 111 | class OldStyleClass: 112 | def foo(self, x): 113 | # type: (Any) -> Any 114 | return x 115 | 116 | def i_care_about_whats_happening(y, z): 117 | # type: (Any, Any) -> Any 118 | print_int(y) 119 | print(z) 120 | return FooReturn() 121 | 122 | 123 | def takes_different_lists(l): 124 | # type: (Any) -> Any 125 | pass 126 | 127 | 128 | def takes_int_lists(l): 129 | # type: (Any) -> Any 130 | pass 131 | 132 | 133 | def takes_int_float_lists(l): 134 | # type: (Any) -> Any 135 | pass 136 | 137 | 138 | def takes_int_to_str_dict(d): 139 | # type: (Any) -> Any 140 | pass 141 | 142 | 143 | def takes_int_to_multiple_val_dict(d): 144 | # type: (Any) -> Any 145 | pass 146 | 147 | 148 | def recursive_dict(d): 149 | # type: (Any) -> Any 150 | pass 151 | 152 | 153 | def empty_then_not_dict(d): 154 | # type: (Any) -> Any 155 | return d 156 | 157 | 158 | def empty_then_not_list(l): 159 | # type: (Any) -> Any 160 | pass 161 | 162 | 163 | def tuple_verify(t): 164 | # type: (Any) -> Any 165 | return t 166 | 167 | 168 | def problematic_dup(uni, bol): 169 | # type: (Text, bool) -> Tuple[Dict[Text, Union[List, int, Text]],bytes] 170 | return {u"foo": [], u"bart": u'ads', u"bax": 23}, b'str' 171 | 172 | 173 | def two_dict_comprehensions(): 174 | # type: () -> Dict[int, Dict[Tuple[int, int], int]] 175 | d = {1: {1: 2}} 176 | return { 177 | i: { 178 | (i, k): l 179 | for k, l in j.items() 180 | } 181 | for i, j in d.items() 182 | } 183 | 184 | 185 | class TestBaseClass(unittest.TestCase): 186 | 187 | def setUp(self): 188 | # type: () -> None 189 | super(TestBaseClass, self).setUp() 190 | # Stats in the same format as the generated JSON. 191 | self.stats = [] # type: List[collect_types.FunctionData] 192 | 193 | def tearDown(self): 194 | # type: () -> None 195 | collect_types.stop_types_collection() 196 | 197 | def load_stats(self): 198 | # type: () -> None 199 | self.stats = json.loads(collect_types.dumps_stats()) 200 | 201 | @contextlib.contextmanager 202 | def collecting_types(self): 203 | # type: () -> Iterator[None] 204 | collect_types.collected_args = {} 205 | collect_types.collected_signatures = {} 206 | collect_types.num_samples = {} 207 | collect_types.sampling_counters = {} 208 | collect_types.call_pending = set() 209 | collect_types.start() 210 | yield None 211 | collect_types.stop() 212 | self.load_stats() 213 | 214 | def assert_type_comments(self, func_name, comments): 215 | # type: (str, List[str]) -> None 216 | """Assert that we generated expected comment for the func_name function in self.stats""" 217 | stat_items = [item for item in self.stats if item.get('func_name') == func_name] 218 | if not comments and not stat_items: 219 | # If we expect no comments, it's okay if nothing was collected. 220 | return 221 | assert len(stat_items) == 1 222 | item = stat_items[0] 223 | if set(item['type_comments']) != set(comments): 224 | print('Actual:') 225 | for comment in sorted(item['type_comments']): 226 | print(' ' + comment) 227 | print('Expected:') 228 | for comment in sorted(comments): 229 | print(' ' + comment) 230 | assert set(item['type_comments']) == set(comments) 231 | assert len(item['type_comments']) == len(comments) 232 | assert os.path.join(collect_types.TOP_DIR, item['path']) == __file__ 233 | 234 | 235 | class TestCollectTypes(TestBaseClass): 236 | 237 | def setUp(self): 238 | # type: () -> None 239 | super(TestCollectTypes, self).setUp() 240 | collect_types.init_types_collection() 241 | 242 | # following type annotations are intentionally use Any, 243 | # because we are testing runtime type collection 244 | 245 | def foo(self, int_arg, list_arg): 246 | # type: (Any, Any) -> None 247 | """foo""" 248 | self.bar(int_arg, list_arg) 249 | 250 | def bar(self, int_arg, list_arg): 251 | # type: (Any, Any) -> Any 252 | """bar""" 253 | return len(self.baz(list_arg)) + int_arg 254 | 255 | def baz(self, list_arg): 256 | # type: (Any) -> Any 257 | """baz""" 258 | return set([int(s) for s in list_arg]) 259 | 260 | def test_type_collection_on_main_thread(self): 261 | # type: () -> None 262 | with self.collecting_types(): 263 | self.foo(2, ['1', '2']) 264 | self.assert_type_comments('TestCollectTypes.foo', ['(int, List[str]) -> None']) 265 | self.assert_type_comments('TestCollectTypes.bar', ['(int, List[str]) -> int']) 266 | self.assert_type_comments('TestCollectTypes.baz', ['(List[str]) -> Set[int]']) 267 | 268 | def bar_another_thread(self, int_arg, list_arg): 269 | # type: (Any, Any) -> Any 270 | """bar""" 271 | return len(self.baz_another_thread(list_arg)) + int_arg 272 | 273 | def baz_another_thread(self, list_arg): 274 | # type: (Any) -> Any 275 | """baz""" 276 | return set([int(s) for s in list_arg]) 277 | 278 | def test_type_collection_on_another_thread(self): 279 | # type: () -> None 280 | with self.collecting_types(): 281 | t = Thread(target=self.bar_another_thread, args=(100, ['1', '2', '3'],)) 282 | t.start() 283 | t.join() 284 | self.assert_type_comments('TestCollectTypes.baz_another_thread', 285 | ['(List[str]) -> Set[int]']) 286 | 287 | def test_run_a_bunch_of_tests(self): 288 | # type: () -> None 289 | with self.collecting_types(): 290 | to = FooObject() 291 | wc = WorkerClass(42, to) 292 | s = sched.scheduler(time.time, time.sleep) 293 | event_source = EventfulHappenings() 294 | s.enter(.001, 1, wc.do_work, ([52, 'foo,', 32], FooNamedTuple('ab', 97))) 295 | s.enter(.002, 1, wc.do_work, ([52, 32], FooNamedTuple('bc', 98))) 296 | s.enter(.003, 1, wc.do_work_clsmthd, (52, FooNamedTuple('de', 99))) 297 | s.enter(.004, 1, event_source.add_handler, (i_care_about_whats_happening,)) 298 | s.enter(.005, 1, event_source.add_handler, (lambda a, b: print_int(a),)) 299 | s.enter(.006, 1, event_source.something_happened, (1, 'tada')) 300 | s.run() 301 | 302 | takes_different_lists([42, 'as', 323, 'a']) 303 | takes_int_lists([42, 323, 3231]) 304 | takes_int_float_lists([42, 323.2132, 3231]) 305 | takes_int_to_str_dict({2: 'a', 4: 'd'}) 306 | takes_int_to_multiple_val_dict({3: 'a', 4: None, 5: 232}) 307 | recursive_dict({3: {3: 'd'}, 4: {3: 'd'}}) 308 | 309 | empty_then_not_dict({}) 310 | empty_then_not_dict({3: {3: 'd'}, 4: {3: 'd'}}) 311 | empty_then_not_list([]) 312 | empty_then_not_list([1, 2]) 313 | empty_then_not_list([1, 2]) 314 | tuple_verify((1, '4')) 315 | tuple_verify((1, '4')) 316 | 317 | problematic_dup(u'ha', False) 318 | problematic_dup(u'ha', False) 319 | 320 | OldStyleClass().foo(10) 321 | 322 | discard(FooObject.FooNested()) 323 | 324 | # TODO(svorobev): add checks for the rest of the functions 325 | # print_int, 326 | self.assert_type_comments( 327 | 'WorkerClass.__init__', 328 | ['(int, pyannotate_runtime.tests.test_collect_types.FooObject) -> None']) 329 | self.assert_type_comments( 330 | 'do_work_clsmthd', 331 | ['(int, pyannotate_runtime.tests.test_collect_types.FooNamedTuple) -> EOFError']) 332 | self.assert_type_comments('OldStyleClass.foo', ['(int) -> int']) 333 | 334 | # Need __qualname__ to get this right 335 | if sys.version_info >= (3, 3): 336 | self.assert_type_comments( 337 | 'discard', 338 | ['(pyannotate_runtime.tests.test_collect_types:FooObject.FooNested) -> None']) 339 | 340 | # TODO: that could be better 341 | self.assert_type_comments('takes_different_lists', ['(List[Union[int, str]]) -> None']) 342 | 343 | # TODO: that should work 344 | # self.assert_type_comment('empty_then_not_dict', 345 | # '(Dict[int, Dict[int, str]]) -> Dict[int, Dict[int, str]]') 346 | self.assert_type_comments('empty_then_not_list', ['(List[int]) -> None', 347 | '(List) -> None']) 348 | if PY2: 349 | self.assert_type_comments( 350 | 'problematic_dup', 351 | ['(unicode, bool) -> Tuple[Dict[unicode, Union[List, int, unicode]], str]']) 352 | else: 353 | self.assert_type_comments( 354 | 'problematic_dup', 355 | ['(str, bool) -> Tuple[Dict[str, Union[List, int, str]], bytes]']) 356 | 357 | def test_two_signatures(self): 358 | # type: () -> None 359 | 360 | def identity(x): 361 | # type: (Any) -> Any 362 | return x 363 | 364 | with self.collecting_types(): 365 | identity(1) 366 | identity('x') 367 | self.assert_type_comments('identity', ['(int) -> int', '(str) -> str']) 368 | 369 | def test_many_signatures(self): 370 | # type: () -> None 371 | 372 | def identity2(x): 373 | # type: (Any) -> Any 374 | return x 375 | 376 | with self.collecting_types(): 377 | for x in 1, 'x', 2, 'y', slice(1), 1.1, None, False, bytearray(), (), [], set(): 378 | for _ in range(50): 379 | identity2(x) 380 | # We collect at most 8 distinct signatures. 381 | self.assert_type_comments('identity2', ['(int) -> int', 382 | '(str) -> str', 383 | '(slice) -> slice', 384 | '(float) -> float', 385 | '(None) -> None', 386 | '(bool) -> bool', 387 | '(bytearray) -> bytearray', 388 | '(Tuple[]) -> Tuple[]']) 389 | 390 | def test_default_args(self): 391 | # type: () -> None 392 | 393 | def func_default(x=0, y=None): 394 | # type: (Any, Any) -> Any 395 | return x 396 | 397 | with self.collecting_types(): 398 | func_default() 399 | func_default('') 400 | func_default(1.1, True) 401 | self.assert_type_comments('func_default', ['(int, None) -> int', 402 | '(str, None) -> str', 403 | '(float, bool) -> float']) 404 | 405 | def test_keyword_args(self): 406 | # type: () -> None 407 | 408 | def func_kw(x, y): 409 | # type: (Any, Any) -> Any 410 | return x 411 | 412 | with self.collecting_types(): 413 | func_kw(y=1, x='') 414 | func_kw(**{'x': 1.1, 'y': None}) 415 | self.assert_type_comments('func_kw', ['(str, int) -> str', 416 | '(float, None) -> float']) 417 | 418 | def test_no_return(self): 419 | # type: () -> None 420 | 421 | def func_always_fail(x): 422 | # type: (Any) -> Any 423 | raise ValueError 424 | 425 | def func_sometimes_fail(x): 426 | # type: (Any) -> Any 427 | if x == 0: 428 | raise RuntimeError 429 | return x 430 | 431 | with self.collecting_types(): 432 | try: 433 | func_always_fail(1) 434 | except Exception: 435 | pass 436 | try: 437 | func_always_fail('') 438 | except Exception: 439 | pass 440 | try: 441 | func_always_fail(1) 442 | except Exception: 443 | pass 444 | try: 445 | func_sometimes_fail(0) 446 | except Exception: 447 | pass 448 | func_sometimes_fail('') 449 | try: 450 | func_sometimes_fail(0) 451 | except Exception: 452 | pass 453 | self.assert_type_comments('func_always_fail', ['(int) -> pyannotate_runtime.collect_types.NoReturnType', 454 | '(str) -> pyannotate_runtime.collect_types.NoReturnType']) 455 | self.assert_type_comments('func_sometimes_fail', ['(int) -> pyannotate_runtime.collect_types.NoReturnType', 456 | '(str) -> str']) 457 | 458 | def test_only_return(self): 459 | # type: () -> None 460 | 461 | def only_return(x): 462 | # type: (int) -> str 463 | collect_types.start() 464 | return '' 465 | 466 | only_return(1) 467 | collect_types.stop() 468 | self.load_stats() 469 | # No entry is stored if we only have a return event with no matching call. 470 | self.assert_type_comments('only_return', []) 471 | 472 | def test_callee_star_args(self): 473 | # type: () -> None 474 | 475 | def callee_star_args(x, *y): 476 | # type: (Any, *Any) -> Any 477 | return 0 478 | 479 | with self.collecting_types(): 480 | callee_star_args(0) 481 | callee_star_args(1, '') 482 | callee_star_args(slice(1), 1.1, True) 483 | callee_star_args(*(False, 1.1, '')) 484 | self.assert_type_comments('callee_star_args', ['(int) -> int', 485 | '(int, *str) -> int', 486 | '(slice, *Union[bool, float]) -> int', 487 | '(bool, *Union[float, str]) -> int']) 488 | 489 | def test_caller_star_args(self): 490 | # type: () -> None 491 | 492 | def caller_star_args(x, y=None): 493 | # type: (Any, Any) -> Any 494 | return 0 495 | 496 | with self.collecting_types(): 497 | caller_star_args(*(1,)) 498 | caller_star_args(*('', 1.1)) 499 | self.assert_type_comments('caller_star_args', ['(int, None) -> int', 500 | '(str, float) -> int']) 501 | 502 | def test_star_star_args(self): 503 | # type: () -> None 504 | 505 | def star_star_args(x, **kw): 506 | # type: (Any, **Any) -> Any 507 | return 0 508 | 509 | with self.collecting_types(): 510 | star_star_args(1, y='', z=True) 511 | star_star_args(**{'x': True, 'a': 1.1}) 512 | self.assert_type_comments('star_star_args', ['(int) -> int', 513 | '(bool) -> int']) 514 | 515 | def test_fully_qualified_type_name_with_sub_package(self): 516 | # type: () -> None 517 | 518 | def identity_qualified(x): 519 | # type: (Any) -> Any 520 | return x 521 | 522 | with self.collecting_types(): 523 | identity_qualified(collect_types.TentativeType()) 524 | self.assert_type_comments( 525 | 'identity_qualified', 526 | ['(pyannotate_runtime.collect_types.TentativeType) -> ' 527 | 'pyannotate_runtime.collect_types.TentativeType']) 528 | 529 | def test_recursive_function(self): 530 | # type: () -> None 531 | 532 | def recurse(x): 533 | # type: (Any) -> Any 534 | if len(x) == 0: 535 | return 1.1 536 | else: 537 | recurse(x[1:]) 538 | return x[0] 539 | 540 | with self.collecting_types(): 541 | recurse((1, '', True)) 542 | self.assert_type_comments( 543 | 'recurse', 544 | ['(Tuple[]) -> float', 545 | '(Tuple[bool]) -> pyannotate_runtime.collect_types.UnknownType', 546 | '(Tuple[str, bool]) -> pyannotate_runtime.collect_types.UnknownType', 547 | '(Tuple[int, str, bool]) -> pyannotate_runtime.collect_types.UnknownType']) 548 | 549 | def test_recursive_function_2(self): 550 | # type: () -> None 551 | 552 | def recurse(x): 553 | # type: (Any) -> Any 554 | if x == 0: 555 | recurse('') 556 | recurse(1.1) 557 | return False 558 | else: 559 | return x 560 | 561 | with self.collecting_types(): 562 | # The return event for the initial call is mismatched because of 563 | # the recursive calls, so we'll have to drop the return type. 564 | recurse(0) 565 | self.assert_type_comments( 566 | 'recurse', 567 | ['(str) -> str', 568 | '(float) -> float', 569 | '(int) -> pyannotate_runtime.collect_types.UnknownType']) 570 | 571 | def test_ignoring_c_calls(self): 572 | # type: () -> None 573 | 574 | def func(x): 575 | # type: (Any) -> Any 576 | a = [1] 577 | # Each of these generates a c_call/c_return event pair. 578 | y = len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a) 579 | y = len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a), len(a) 580 | str(y) 581 | return x 582 | 583 | with self.collecting_types(): 584 | func(1) 585 | func('') 586 | self.assert_type_comments('func', ['(int) -> int', 587 | '(str) -> str']) 588 | 589 | def test_no_crash_on_nested_dict_comps(self): 590 | # type: () -> None 591 | with self.collecting_types(): 592 | two_dict_comprehensions() 593 | self.assert_type_comments('two_dict_comprehensions', 594 | ['() -> Dict[int, Dict[Tuple[int, int], int]]']) 595 | 596 | def test_skip_lambda(self): 597 | # type: () -> None 598 | with self.collecting_types(): 599 | (lambda: None)() 600 | (lambda x: x)(0) 601 | (lambda x, y: x+y)(0, 0) 602 | assert self.stats == [] 603 | 604 | def test_unknown_module_types(self): 605 | # type: () -> None 606 | def func_with_unknown_module_types(c): 607 | # type: (Any) -> Any 608 | return c 609 | 610 | with self.collecting_types(): 611 | ns = { 612 | '__name__': '' 613 | } # type: Dict[str, Any] 614 | exec('class C(object): pass', ns) 615 | 616 | func_with_unknown_module_types(ns['C']()) 617 | 618 | self.assert_type_comments('func_with_unknown_module_types', ['(C) -> C']) 619 | 620 | def test_yield_basic(self): 621 | # type: () -> None 622 | def gen(n, a): 623 | for i in range(n): 624 | yield a 625 | 626 | with self.collecting_types(): 627 | list(gen(10, 'x')) 628 | 629 | self.assert_type_comments('gen', ['(int, str) -> Iterator[str]']) 630 | 631 | def test_yield_various(self): 632 | # type: () -> None 633 | def gen(n, a, b): 634 | for i in range(n): 635 | yield a 636 | yield b 637 | 638 | with self.collecting_types(): 639 | list(gen(10, 'x', 1)) 640 | list(gen(0, 0, 0)) 641 | 642 | # TODO: This should really return Iterator[Union[int, str]] 643 | self.assert_type_comments('gen', ['(int, str, int) -> Iterator[int]', 644 | '(int, str, int) -> Iterator[str]']) 645 | 646 | def test_yield_empty(self): 647 | # type: () -> None 648 | def gen(): 649 | if False: 650 | yield 651 | 652 | with self.collecting_types(): 653 | list(gen()) 654 | 655 | self.assert_type_comments('gen', ['() -> Iterator']) 656 | 657 | 658 | def foo(arg): 659 | # type: (Any) -> Any 660 | return [arg] 661 | 662 | 663 | class TestInitWithFilter(TestBaseClass): 664 | 665 | def always_foo(self, filename): 666 | # type: (Optional[str]) -> Optional[str] 667 | return 'foo.py' 668 | 669 | def always_none(self, filename): 670 | # type: (Optional[str]) -> Optional[str] 671 | return None 672 | 673 | def test_init_with_filter(self): 674 | # type: () -> None 675 | collect_types.init_types_collection(self.always_foo) 676 | with self.collecting_types(): 677 | foo(42) 678 | assert len(self.stats) == 1 679 | assert self.stats[0]['path'] == 'foo.py' 680 | 681 | def test_init_with_none_filter(self): 682 | # type: () -> None 683 | collect_types.init_types_collection(self.always_none) 684 | with self.collecting_types(): 685 | foo(42) 686 | assert self.stats == [] 687 | -------------------------------------------------------------------------------- /pyannotate_tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dropbox/pyannotate/a7a46f394f0ba91a1b5fbf657e2393af542969ae/pyannotate_tools/__init__.py -------------------------------------------------------------------------------- /pyannotate_tools/annotations/__init__.py: -------------------------------------------------------------------------------- 1 | # type: ignore 2 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/__main__.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import argparse 4 | import json 5 | import logging 6 | import os 7 | import sys 8 | 9 | from lib2to3.main import StdoutRefactoringTool 10 | 11 | from typing import Any, Dict, List, Optional 12 | 13 | from pyannotate_tools.annotations.main import generate_annotations_json_string, unify_type_comments 14 | from pyannotate_tools.fixes.fix_annotate_json import FixAnnotateJson 15 | 16 | parser = argparse.ArgumentParser() 17 | parser.add_argument('--type-info', default='type_info.json', metavar="FILE", 18 | help="JSON input file (default type_info.json)") 19 | parser.add_argument('--uses-signature', action='store_true', 20 | help="JSON input uses a signature format") 21 | parser.add_argument('-p', '--print-function', action='store_true', 22 | help="Assume print is a function") 23 | parser.add_argument('-w', '--write', action='store_true', 24 | help="Write output files") 25 | parser.add_argument('-j', '--processes', type=int, default=1, metavar="N", 26 | help="Use N parallel processes (default no parallelism)") 27 | parser.add_argument('--max-line-drift', type=int, default=5, metavar="N", 28 | help="Maximum allowed line drift when inserting annotation" 29 | " (can be useful for custom codecs)") 30 | parser.add_argument('-v', '--verbose', action='store_true', 31 | help="More verbose output") 32 | parser.add_argument('-q', '--quiet', action='store_true', 33 | help="Don't show diffs") 34 | parser.add_argument('-d', '--dump', action='store_true', 35 | help="Dump raw type annotations (filter by files, default all)") 36 | parser.add_argument('-a', '--auto-any', action='store_true', 37 | help="Annotate everything with 'Any', without reading type_info.json") 38 | parser.add_argument('files', nargs='*', metavar="FILE", 39 | help="Files and directories to update with annotations") 40 | parser.add_argument('-s', '--only-simple', action='store_true', 41 | help="Only annotate functions with trivial types") 42 | parser.add_argument('--python-version', action='store', default='2', choices=['2', '3'], 43 | help="Choose annotation style, 2 for Python 2 with comments (the " 44 | "default), 3 for Python 3 with annotation syntax" ) 45 | parser.add_argument('--py2', '-2', action='store_const', dest='python_version', const='2', 46 | help="Annotate for Python 2 with comments (default)") 47 | parser.add_argument('--py3', '-3', action='store_const', dest='python_version', const='3', 48 | help="Annotate for Python 3 with argument and return value annotations") 49 | 50 | 51 | class ModifiedRefactoringTool(StdoutRefactoringTool): 52 | """Class that gives a nicer error message for bad encodings.""" 53 | 54 | def refactor_file(self, filename, write=False, doctests_only=False): 55 | try: 56 | super(ModifiedRefactoringTool, self).refactor_file( 57 | filename, write=write, doctests_only=doctests_only) 58 | except SyntaxError as err: 59 | if str(err).startswith("unknown encoding:"): 60 | self.log_error("Can't parse %s: %s", filename, err) 61 | else: 62 | raise 63 | 64 | 65 | def dump_annotations(type_info, files): 66 | """Dump annotations out of type_info, filtered by files. 67 | 68 | If files is non-empty, only dump items either if the path in the 69 | item matches one of the files exactly, or else if one of the files 70 | is a path prefix of the path. 71 | """ 72 | with open(type_info) as f: 73 | data = json.load(f) 74 | for item in data: 75 | path, line, func_name = item['path'], item['line'], item['func_name'] 76 | if files and path not in files: 77 | for f in files: 78 | if path.startswith(os.path.join(f, '')): 79 | break 80 | else: 81 | continue # Outer loop 82 | print("%s:%d: in %s:" % (path, line, func_name)) 83 | type_comments = item['type_comments'] 84 | signature = unify_type_comments(type_comments) 85 | arg_types = signature['arg_types'] 86 | return_type = signature['return_type'] 87 | print(" # type: (%s) -> %s" % (", ".join(arg_types), return_type)) 88 | 89 | 90 | def main(args_override=None): 91 | # type: (Optional[List[str]]) -> None 92 | 93 | # Parse command line. 94 | args = parser.parse_args(args_override) 95 | if not args.files and not args.dump: 96 | parser.error("At least one file/directory is required") 97 | 98 | if args.python_version not in ('2', '3'): 99 | sys.exit('--python-version must be 2 or 3') 100 | 101 | annotation_style = 'py' + args.python_version 102 | 103 | # Set up logging handler. 104 | level = logging.DEBUG if args.verbose else logging.INFO 105 | logging.basicConfig(format='%(message)s', level=level) 106 | 107 | if args.dump: 108 | dump_annotations(args.type_info, args.files) 109 | return 110 | 111 | if args.auto_any: 112 | fixers = ['pyannotate_tools.fixes.fix_annotate'] 113 | else: 114 | # Produce nice error message if type_info.json not found. 115 | try: 116 | with open(args.type_info) as f: 117 | contents = f.read() 118 | except IOError as err: 119 | sys.exit("Can't open type info file: %s" % err) 120 | 121 | # Run pass 2 with output into a variable. 122 | if args.uses_signature: 123 | data = json.loads(contents) # type: List[Any] 124 | else: 125 | data = generate_annotations_json_string( 126 | args.type_info, 127 | only_simple=args.only_simple) 128 | 129 | # Run pass 3 with input from that variable. 130 | FixAnnotateJson.init_stub_json_from_data(data, args.files[0]) 131 | fixers = ['pyannotate_tools.fixes.fix_annotate_json'] 132 | 133 | flags = {'print_function': args.print_function, 134 | 'annotation_style': annotation_style} 135 | rt = ModifiedRefactoringTool( 136 | fixers=fixers, 137 | options=flags, 138 | explicit=fixers, 139 | nobackups=True, 140 | show_diffs=not args.quiet) 141 | if not rt.errors: 142 | with FixAnnotateJson.max_line_drift_set(args.max_line_drift): 143 | rt.refactor(args.files, write=args.write, num_processes=args.processes) 144 | if args.processes == 1: 145 | rt.summarize() 146 | else: 147 | logging.info("(In multi-process per-file warnings are lost)") 148 | if not args.write: 149 | logging.info("NOTE: this was a dry run; use -w to write files") 150 | 151 | 152 | if __name__ == '__main__': 153 | main() 154 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/infer.py: -------------------------------------------------------------------------------- 1 | """Infer an annotation from a set of concrete runtime type signatures. 2 | 3 | The main entry point is 'infer_annotation'. 4 | """ 5 | 6 | from typing import Dict, Iterable, List, Optional, Set, Tuple 7 | 8 | from pyannotate_tools.annotations.parse import parse_type_comment 9 | from pyannotate_tools.annotations.types import ( 10 | AbstractType, 11 | AnyType, 12 | ARG_POS, 13 | Argument, 14 | ClassType, 15 | is_optional, 16 | TupleType, 17 | UnionType, 18 | NoReturnType, 19 | ) 20 | 21 | IGNORED_ITEMS = { 22 | 'unittest.mock.Mock', 23 | 'unittest.mock.MagicMock', 24 | 'mock.mock.Mock', 25 | 'mock.mock.MagicMock', 26 | } 27 | 28 | class InferError(Exception): 29 | """Raised if we can't infer a signature for some reason.""" 30 | 31 | 32 | def infer_annotation(type_comments): 33 | # type: (List[str]) -> Tuple[List[Argument], AbstractType] 34 | """Given some type comments, return a single inferred signature. 35 | 36 | Args: 37 | type_comments: Strings of form '(arg1, ... argN) -> ret' 38 | 39 | Returns: Tuple of (argument types and kinds, return type). 40 | """ 41 | assert type_comments 42 | args = {} # type: Dict[int, Set[Argument]] 43 | returns = set() 44 | for comment in type_comments: 45 | arg_types, return_type = parse_type_comment(comment) 46 | for i, arg_type in enumerate(arg_types): 47 | args.setdefault(i, set()).add(arg_type) 48 | returns.add(return_type) 49 | combined_args = [] 50 | for i in sorted(args): 51 | arg_infos = list(args[i]) 52 | kind = argument_kind(arg_infos) 53 | if kind is None: 54 | raise InferError('Ambiguous argument kinds:\n' + '\n'.join(type_comments)) 55 | types = [arg.type for arg in arg_infos] 56 | combined = combine_types(types) 57 | if str(combined) == 'None': 58 | # It's very rare for an argument to actually be typed `None`, more likely than 59 | # not we simply don't have any data points for this argument. 60 | combined = UnionType([ClassType('None'), AnyType()]) 61 | if kind != ARG_POS and (len(str(combined)) > 120 or isinstance(combined, UnionType)): 62 | # Avoid some noise. 63 | combined = AnyType() 64 | combined_args.append(Argument(combined, kind)) 65 | combined_return = combine_types(returns) 66 | return combined_args, combined_return 67 | 68 | 69 | def argument_kind(args): 70 | # type: (List[Argument]) -> Optional[str] 71 | """Return the kind of an argument, based on one or more descriptions of the argument. 72 | 73 | Return None if every item does not have the same kind. 74 | """ 75 | kinds = set(arg.kind for arg in args) 76 | if len(kinds) != 1: 77 | return None 78 | return kinds.pop() 79 | 80 | 81 | def combine_types(types): 82 | # type: (Iterable[AbstractType]) -> AbstractType 83 | """Given some types, return a combined and simplified type. 84 | 85 | For example, if given 'int' and 'List[int]', return Union[int, List[int]]. If given 86 | 'int' and 'int', return just 'int'. 87 | """ 88 | items = simplify_types(types) 89 | if len(items) == 1: 90 | return items[0] 91 | else: 92 | return UnionType(items) 93 | 94 | 95 | def simplify_types(types): 96 | # type: (Iterable[AbstractType]) -> List[AbstractType] 97 | """Given some types, give simplified types representing the union of types.""" 98 | flattened = flatten_types(types) 99 | items = filter_ignored_items(flattened) 100 | items = [simplify_recursive(item) for item in items] 101 | items = merge_items(items) 102 | items = dedupe_types(items) 103 | # We have to remove reundant items after everything has been simplified and 104 | # merged as this simplification may be what makes items redundant. 105 | items = remove_redundant_items(items) 106 | if len(items) > 3: 107 | return [AnyType()] 108 | else: 109 | return items 110 | 111 | 112 | def simplify_recursive(typ): 113 | # type: (AbstractType) -> AbstractType 114 | """Simplify all components of a type.""" 115 | if isinstance(typ, UnionType): 116 | return combine_types(typ.items) 117 | elif isinstance(typ, ClassType): 118 | simplified = ClassType(typ.name, [simplify_recursive(arg) for arg in typ.args]) 119 | args = simplified.args 120 | if (simplified.name == 'Dict' and len(args) == 2 121 | and isinstance(args[0], ClassType) and args[0].name in ('str', 'Text') 122 | and isinstance(args[1], UnionType) and not is_optional(args[1])): 123 | # Looks like a potential case for TypedDict, which we don't properly support yet. 124 | return ClassType('Dict', [args[0], AnyType()]) 125 | return simplified 126 | elif isinstance(typ, TupleType): 127 | return TupleType([simplify_recursive(item) for item in typ.items]) 128 | return typ 129 | 130 | 131 | def flatten_types(types): 132 | # type: (Iterable[AbstractType]) -> List[AbstractType] 133 | flattened = [] 134 | for item in types: 135 | if not isinstance(item, UnionType): 136 | flattened.append(item) 137 | else: 138 | flattened.extend(flatten_types(item.items)) 139 | return flattened 140 | 141 | 142 | def dedupe_types(types): 143 | # type: (Iterable[AbstractType]) -> List[AbstractType] 144 | return sorted(set(types), key=lambda t: str(t)) 145 | 146 | def filter_ignored_items(items): 147 | # type: (List[AbstractType]) -> List[AbstractType] 148 | result = [item for item in items 149 | if not isinstance(item, ClassType) or 150 | item.name not in IGNORED_ITEMS] 151 | return result or [AnyType()] 152 | 153 | def remove_redundant_items(items): 154 | # type: (List[AbstractType]) -> List[AbstractType] 155 | """Filter out redundant union items.""" 156 | result = [] 157 | for item in items: 158 | for other in items: 159 | if item is not other and is_redundant_union_item(item, other): 160 | break 161 | else: 162 | result.append(item) 163 | return result 164 | 165 | 166 | def is_redundant_union_item(first, other): 167 | # type: (AbstractType, AbstractType) -> bool 168 | """If union has both items, is the first one redundant? 169 | 170 | For example, if first is 'str' and the other is 'Text', return True. 171 | 172 | If items are equal, return False. 173 | """ 174 | if isinstance(first, ClassType) and isinstance(other, ClassType): 175 | if first.name == 'str' and other.name == 'Text': 176 | return True 177 | elif first.name == 'bool' and other.name == 'int': 178 | return True 179 | elif first.name == 'int' and other.name == 'float': 180 | return True 181 | elif (first.name in ('List', 'Dict', 'Set') and 182 | other.name == first.name): 183 | if not first.args and other.args: 184 | return True 185 | elif len(first.args) == len(other.args) and first.args: 186 | result = all(first_arg == other_arg or other_arg == AnyType() 187 | for first_arg, other_arg 188 | in zip(first.args, other.args)) 189 | return result 190 | 191 | return False 192 | 193 | 194 | def merge_items(items): 195 | # type: (List[AbstractType]) -> List[AbstractType] 196 | """Merge union items that can be merged.""" 197 | result = [] 198 | while items: 199 | item = items.pop() 200 | merged = None 201 | for i, other in enumerate(items): 202 | merged = merged_type(item, other) 203 | if merged: 204 | break 205 | if merged: 206 | del items[i] 207 | items.append(merged) 208 | else: 209 | result.append(item) 210 | return list(reversed(result)) 211 | 212 | 213 | def merged_type(t, s): 214 | # type: (AbstractType, AbstractType) -> Optional[AbstractType] 215 | """Return merged type if two items can be merged in to a different, more general type. 216 | 217 | Return None if merging is not possible. 218 | """ 219 | if isinstance(t, TupleType) and isinstance(s, TupleType): 220 | if len(t.items) == len(s.items): 221 | return TupleType([combine_types([ti, si]) for ti, si in zip(t.items, s.items)]) 222 | all_items = t.items + s.items 223 | if all_items and all(item == all_items[0] for item in all_items[1:]): 224 | # Merge multiple compatible fixed-length tuples into a variable-length tuple type. 225 | return ClassType('Tuple', [all_items[0]]) 226 | elif (isinstance(t, TupleType) and isinstance(s, ClassType) and s.name == 'Tuple' 227 | and len(s.args) == 1): 228 | if all(item == s.args[0] for item in t.items): 229 | # Merge fixed-length tuple and variable-length tuple. 230 | return s 231 | elif isinstance(s, TupleType) and isinstance(t, ClassType) and t.name == 'Tuple': 232 | return merged_type(s, t) 233 | elif isinstance(s, NoReturnType): 234 | return t 235 | elif isinstance(t, NoReturnType): 236 | return s 237 | elif isinstance(s, AnyType): 238 | # This seems to be usually desirable, since Anys tend to come from unknown types. 239 | return t 240 | elif isinstance(t, AnyType): 241 | # Similar to above. 242 | return s 243 | return None 244 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/main.py: -------------------------------------------------------------------------------- 1 | """Main entry point to mypy annotation inference utility.""" 2 | 3 | import json 4 | 5 | from typing import List 6 | from mypy_extensions import TypedDict 7 | 8 | from pyannotate_tools.annotations.types import ARG_STAR, ARG_STARSTAR 9 | from pyannotate_tools.annotations.infer import infer_annotation 10 | from pyannotate_tools.annotations.parse import parse_json 11 | 12 | 13 | # Schema of a function signature in the output 14 | Signature = TypedDict('Signature', {'arg_types': List[str], 15 | 'return_type': str}) 16 | 17 | # Schema of a function in the output 18 | FunctionData = TypedDict('FunctionData', {'path': str, 19 | 'line': int, 20 | 'func_name': str, 21 | 'signature': Signature, 22 | 'samples': int}) 23 | SIMPLE_TYPES = {'None', 'int', 'float', 'str', 'bytes', 'bool'} 24 | 25 | def unify_type_comments(type_comments): 26 | # type: (List[str]) -> Signature 27 | arg_types, return_type = infer_annotation(type_comments) 28 | arg_strs = [] 29 | for arg, kind in arg_types: 30 | arg_str = str(arg) 31 | if kind == ARG_STAR: 32 | arg_str = '*%s' % arg_str 33 | elif kind == ARG_STARSTAR: 34 | arg_str = '**%s' % arg_str 35 | arg_strs.append(arg_str) 36 | return { 37 | 'arg_types': arg_strs, 38 | 'return_type': str(return_type), 39 | } 40 | 41 | 42 | def is_signature_simple(signature): 43 | # type: (Signature) -> bool 44 | return (all(x.lstrip('*') in SIMPLE_TYPES for x in signature['arg_types']) and 45 | signature['return_type'] in SIMPLE_TYPES) 46 | 47 | 48 | def generate_annotations_json_string(source_path, only_simple=False): 49 | # type: (str, bool) -> List[FunctionData] 50 | """Produce annotation data JSON file from a JSON file with runtime-collected types. 51 | 52 | Data formats: 53 | 54 | * The source JSON is a list of pyannotate_tools.annotations.parse.RawEntry items. 55 | * The output JSON is a list of FunctionData items. 56 | """ 57 | items = parse_json(source_path) 58 | results = [] 59 | for item in items: 60 | signature = unify_type_comments(item.type_comments) 61 | if is_signature_simple(signature) or not only_simple: 62 | data = { 63 | 'path': item.path, 64 | 'line': item.line, 65 | 'func_name': item.func_name, 66 | 'signature': signature, 67 | 'samples': item.samples 68 | } # type: FunctionData 69 | results.append(data) 70 | return results 71 | 72 | 73 | def generate_annotations_json(source_path, target_path, only_simple=False): 74 | # type: (str, str, bool) -> None 75 | """Like generate_annotations_json_string() but writes JSON to a file.""" 76 | results = generate_annotations_json_string(source_path, only_simple=only_simple) 77 | with open(target_path, 'w') as f: 78 | json.dump(results, f, sort_keys=True, indent=4) 79 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/parse.py: -------------------------------------------------------------------------------- 1 | """Parse type annotations collected at runtime by collect_types and dumped as JSON. 2 | 3 | Parse JSON data and also parse type comment strings into type objects. 4 | 5 | The collect_types tool is in pyannotate_runtime/collect_types.py. 6 | """ 7 | 8 | import json 9 | import re 10 | import sys 11 | 12 | from typing import Any, List, Mapping, Set, Tuple 13 | try: 14 | from typing import Text 15 | except ImportError: 16 | # In Python 3.5.1 stdlib, typing.py does not define Text 17 | Text = str # type: ignore 18 | from mypy_extensions import NoReturn, TypedDict 19 | 20 | from pyannotate_tools.annotations.types import ( 21 | AbstractType, 22 | AnyType, 23 | ARG_POS, 24 | ARG_STAR, 25 | ARG_STARSTAR, 26 | Argument, 27 | ClassType, 28 | TupleType, 29 | UnionType, 30 | NoReturnType, 31 | ) 32 | 33 | PY2 = sys.version_info < (3,) 34 | 35 | 36 | # Rules for replacing some type names that aren't valid Python names or that 37 | # are otherwise invalid. 38 | TYPE_FIXUPS = { 39 | # The dictionary-* names come from Python 2 `__class__.__name__` values 40 | # from `dict.iterkeys()`, etc. Python 3 uses valid names. 41 | 'dictionary-keyiterator': 'Iterator', 42 | 'dictionary-valueiterator': 'Iterator', 43 | 'dictionary-itemiterator': 'Iterator', 44 | 'pyannotate_runtime.collect_types.UnknownType': 'Any', 45 | 'pyannotate_runtime.collect_types.NoReturnType': 'mypy_extensions.NoReturn', 46 | 'function': 'Callable', 47 | 'functools.partial': 'Callable', 48 | 'long': 'int', 49 | 'unicode': 'Text', 50 | 'generator': 'Iterator', 51 | 'listiterator': 'Iterator', 52 | 'instancemethod': 'Callable', 53 | 'itertools.imap': 'Iterator', 54 | 'operator.methodcaller': 'Callable', 55 | 'method': 'Callable', 56 | 'method-wrapper': 'Callable', 57 | 'mappingproxy': 'Mapping', 58 | 'file': 'IO[bytes]', 59 | 'instance': 'Any', 60 | 'collections.defaultdict': 'Dict', 61 | } 62 | 63 | 64 | # Input JSON data entry 65 | RawEntry = TypedDict('RawEntry', {'path': Text, 66 | 'line': int, 67 | 'func_name': Text, 68 | 'type_comments': List[Text], 69 | 'samples': int}) 70 | 71 | 72 | class FunctionInfo(object): 73 | """Deserialized raw runtime information for a single function (based on RawEntry)""" 74 | 75 | def __init__(self, path, line, func_name, type_comments, samples): 76 | # type: (str, int, str, List[str], int) -> None 77 | self.path = path 78 | self.line = line 79 | self.func_name = func_name 80 | self.type_comments = type_comments 81 | self.samples = samples 82 | 83 | 84 | class ParseError(Exception): 85 | """Raised on any type comment parse error. 86 | 87 | The 'comment' attribute contains the comment that produced the error. 88 | """ 89 | 90 | def __init__(self, comment): 91 | # type: (str) -> None 92 | super(ParseError, self).__init__('Invalid type comment: %s' % comment) 93 | self.comment = comment 94 | 95 | 96 | def parse_json(path): 97 | # type: (str) -> List[FunctionInfo] 98 | """Deserialize a JSON file containing runtime collected types. 99 | 100 | The input JSON is expected to to have a list of RawEntry items. 101 | """ 102 | with open(path) as f: 103 | data = json.load(f) # type: List[RawEntry] 104 | result = [] 105 | 106 | def assert_type(value, typ): 107 | # type: (object, type) -> None 108 | assert isinstance(value, typ), '%s: Unexpected type %r' % (path, type(value).__name__) 109 | 110 | def assert_dict_item(dictionary, key, typ): 111 | # type: (Mapping[Any, Any], str, type) -> None 112 | assert key in dictionary, '%s: Missing dictionary key %r' % (path, key) 113 | value = dictionary[key] 114 | assert isinstance(value, typ), '%s: Unexpected type %r for key %r' % ( 115 | path, type(value).__name__, key) 116 | 117 | assert_type(data, list) 118 | for item in data: 119 | assert_type(item, dict) 120 | assert_dict_item(item, 'path', Text) 121 | assert_dict_item(item, 'line', int) 122 | assert_dict_item(item, 'func_name', Text) 123 | assert_dict_item(item, 'type_comments', list) 124 | for comment in item['type_comments']: 125 | assert_type(comment, Text) 126 | assert_type(item['samples'], int) 127 | info = FunctionInfo(encode(item['path']), 128 | item['line'], 129 | encode(item['func_name']), 130 | [encode(comment) for comment in item['type_comments']], 131 | item['samples']) 132 | result.append(info) 133 | return result 134 | 135 | 136 | class Token(object): 137 | """Abstract base class for tokens used for parsing type comments""" 138 | text = '' 139 | 140 | 141 | class DottedName(Token): 142 | """An identifier token, such as 'List', 'int' or 'package.name'""" 143 | 144 | def __init__(self, text): 145 | # type: (str) -> None 146 | self.text = text 147 | 148 | def __repr__(self): 149 | # type: () -> str 150 | return 'DottedName(%s)' % self.text 151 | 152 | 153 | class Separator(Token): 154 | """A separator or punctuator token such as '(', '[' or '->'""" 155 | 156 | def __init__(self, text): 157 | # type: (str) -> None 158 | self.text = text 159 | 160 | def __repr__(self): 161 | # type: () -> str 162 | return self.text 163 | 164 | 165 | class End(Token): 166 | """A token representing the end of a type comment""" 167 | 168 | def __repr__(self): 169 | # type: () -> str 170 | return 'End()' 171 | 172 | 173 | def tokenize(s): 174 | # type: (str) -> List[Token] 175 | """Translate a type comment into a list of tokens.""" 176 | original = s 177 | tokens = [] # type: List[Token] 178 | while True: 179 | if not s: 180 | tokens.append(End()) 181 | return tokens 182 | elif s[0] == ' ': 183 | s = s[1:] 184 | elif s[0] in '()[],*': 185 | tokens.append(Separator(s[0])) 186 | s = s[1:] 187 | elif s[:2] == '->': 188 | tokens.append(Separator('->')) 189 | s = s[2:] 190 | else: 191 | m = re.match(r'[-\w]+(\s*(\.|:)\s*[-/\w]*)*', s) 192 | if not m: 193 | raise ParseError(original) 194 | fullname = m.group(0) 195 | fullname = fullname.replace(' ', '') 196 | if fullname in TYPE_FIXUPS: 197 | fullname = TYPE_FIXUPS[fullname] 198 | # pytz creates classes with the name of the timezone being used: 199 | # https://github.com/stub42/pytz/blob/f55399cddbef67c56db1b83e0939ecc1e276cf42/src/pytz/tzfile.py#L120-L123 200 | # This causes pyannotates to crash as it's invalid to have a class 201 | # name with a `/` in it (e.g. "pytz.tzfile.America/Los_Angeles") 202 | if fullname.startswith('pytz.tzfile.'): 203 | fullname = 'datetime.tzinfo' 204 | if '-' in fullname or '/' in fullname: 205 | # Not a valid Python name; there are many places that 206 | # generate these, so we just substitute Any rather 207 | # than crashing. 208 | fullname = 'Any' 209 | tokens.append(DottedName(fullname)) 210 | s = s[len(m.group(0)):] 211 | 212 | 213 | def parse_type_comment(comment): 214 | # type: (str) -> Tuple[List[Argument], AbstractType] 215 | """Parse a type comment of form '(arg1, ..., argN) -> ret'.""" 216 | return Parser(comment).parse() 217 | 218 | 219 | class Parser(object): 220 | """Implementation of the type comment parser""" 221 | 222 | def __init__(self, comment): 223 | # type: (str) -> None 224 | self.comment = comment 225 | self.tokens = tokenize(comment) 226 | self.i = 0 227 | 228 | def parse(self): 229 | # type: () -> Tuple[List[Argument], AbstractType] 230 | self.expect('(') 231 | arg_types = [] # type: List[Argument] 232 | stars_seen = set() # type: Set[str] 233 | while self.lookup() != ')': 234 | if self.lookup() == '*': 235 | self.expect('*') 236 | if self.lookup() == '*': 237 | if '**' in stars_seen: 238 | self.fail() 239 | self.expect('*') 240 | star_star = True 241 | else: 242 | if stars_seen: 243 | self.fail() 244 | star_star = False 245 | arg_type = self.parse_type() 246 | if star_star: 247 | arg_types.append(Argument(arg_type, ARG_STARSTAR)) 248 | stars_seen.add('**') 249 | else: 250 | arg_types.append(Argument(arg_type, ARG_STAR)) 251 | stars_seen.add('*') 252 | else: 253 | if stars_seen: 254 | self.fail() 255 | arg_type = self.parse_type() 256 | arg_types.append(Argument(arg_type, ARG_POS)) 257 | if self.lookup() == ',': 258 | self.expect(',') 259 | elif self.lookup() == ')': 260 | break 261 | self.expect(')') 262 | self.expect('->') 263 | ret_type = self.parse_type() 264 | if not isinstance(self.next(), End): 265 | self.fail() 266 | return arg_types, ret_type 267 | 268 | def parse_type_list(self): 269 | # type: () -> List[AbstractType] 270 | types = [] 271 | while self.lookup() not in (')', ']'): 272 | typ = self.parse_type() 273 | types.append(typ) 274 | if self.lookup() == ',': 275 | self.expect(',') 276 | elif self.lookup() not in (')', ']'): 277 | self.fail() 278 | return types 279 | 280 | def parse_type(self): 281 | # type: () -> AbstractType 282 | t = self.next() 283 | if not isinstance(t, DottedName): 284 | self.fail() 285 | if t.text == 'Any': 286 | return AnyType() 287 | elif t.text == 'mypy_extensions.NoReturn': 288 | return NoReturnType() 289 | elif t.text == 'Tuple': 290 | self.expect('[') 291 | args = self.parse_type_list() 292 | self.expect(']') 293 | return TupleType(args) 294 | elif t.text == 'Union': 295 | self.expect('[') 296 | items = self.parse_type_list() 297 | self.expect(']') 298 | if len(items) == 1: 299 | return items[0] 300 | elif len(items) == 0: 301 | self.fail() 302 | else: 303 | return UnionType(items) 304 | else: 305 | if self.lookup() == '[': 306 | self.expect('[') 307 | args = self.parse_type_list() 308 | self.expect(']') 309 | if t.text == 'Optional' and len(args) == 1: 310 | return UnionType([args[0], ClassType('None')]) 311 | return ClassType(t.text, args) 312 | else: 313 | return ClassType(t.text) 314 | 315 | def expect(self, s): 316 | # type: (str) -> None 317 | if self.tokens[self.i].text != s: 318 | self.fail() 319 | self.i += 1 320 | 321 | def lookup(self): 322 | # type: () -> str 323 | return self.tokens[self.i].text 324 | 325 | def next(self): 326 | # type: () -> Token 327 | token = self.tokens[self.i] 328 | self.i += 1 329 | return token 330 | 331 | def fail(self): 332 | # type: () -> NoReturn 333 | raise ParseError(self.comment) 334 | 335 | 336 | def encode(s): 337 | # type: (Text) -> str 338 | if PY2: 339 | return s.encode('ascii') 340 | else: 341 | return s 342 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/__init__.py: -------------------------------------------------------------------------------- 1 | # type: ignore 2 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/dundermain_test.py: -------------------------------------------------------------------------------- 1 | """Some (nearly) end-to-end testing.""" 2 | 3 | import json 4 | import os 5 | import re 6 | import shutil 7 | import sys 8 | import tempfile 9 | import unittest 10 | 11 | # There seems to be no way to have this work and type-check without an 12 | # explicit version check. :-( 13 | if sys.version_info[0] == 2: 14 | from cStringIO import StringIO 15 | else: 16 | from io import StringIO 17 | 18 | from typing import Iterator, List 19 | 20 | from pyannotate_tools.annotations.__main__ import main as dunder_main 21 | 22 | 23 | class TestDunderMain(unittest.TestCase): 24 | def setUp(self): 25 | # type: () -> None 26 | self.tempdirname = tempfile.mkdtemp() 27 | self.tempfiles = [] # type: List[str] 28 | self.olddir = os.getcwd() 29 | os.chdir(self.tempdirname) 30 | 31 | def tearDown(self): 32 | # type: () -> None 33 | os.chdir(self.olddir) 34 | shutil.rmtree(self.tempdirname) 35 | 36 | def write_file(self, name, data): 37 | # type: (str, str) -> None 38 | self.tempfiles.append(name) 39 | with open(name, 'w') as f: 40 | f.write(data) 41 | 42 | def test_help(self): 43 | # type: () -> None 44 | self.main_test(["--help"], r"^usage:", r"^$", 0) 45 | 46 | def test_preview(self): 47 | # type: () -> None 48 | self.prototype_test(write=False) 49 | 50 | def test_final(self): 51 | # type: () -> None 52 | self.prototype_test(write=True) 53 | with open('gcd.py') as f: 54 | lines = [line.strip() for line in f.readlines()] 55 | assert '# type: (int, int) -> int' in lines 56 | 57 | def test_bad_encoding_message(self): 58 | # type: () -> None 59 | source_text = "# coding: unknownencoding\ndef f():\n pass\n" 60 | self.write_file('gcd.py', source_text) 61 | self.write_file('type_info.json', '[]') 62 | encoding_message = "Can't parse gcd.py: unknown encoding: unknownencoding" 63 | self.main_test(['gcd.py'], 64 | r'\A\Z', 65 | r'\A' + re.escape(encoding_message), 66 | 0) 67 | 68 | def prototype_test(self, write): 69 | # type: (bool) -> None 70 | type_info = [ 71 | { 72 | "path": "gcd.py", 73 | "line": 1, 74 | "func_name": "gcd", 75 | "type_comments": [ 76 | "(int, int) -> int" 77 | ], 78 | "samples": 2 79 | } 80 | ] 81 | source_text = """\ 82 | def gcd(a, b): 83 | while b: 84 | a, b = b, a%b 85 | return a 86 | """ 87 | stdout_expected = """\ 88 | --- gcd.py (original) 89 | +++ gcd.py (refactored) 90 | @@ -1,4 +1,5 @@ 91 | def gcd(a, b): 92 | + # type: (int, int) -> int 93 | while b: 94 | a, b = b, a%b 95 | return a 96 | """ 97 | if not write: 98 | stderr_expected = """\ 99 | Refactored gcd.py 100 | Files that need to be modified: 101 | gcd.py 102 | NOTE: this was a dry run; use -w to write files 103 | """ 104 | else: 105 | stderr_expected = """\ 106 | Refactored gcd.py 107 | Files that were modified: 108 | gcd.py 109 | """ 110 | self.write_file('type_info.json', json.dumps(type_info)) 111 | self.write_file('gcd.py', source_text) 112 | args = ['gcd.py'] 113 | if write: 114 | args.append('-w') 115 | self.main_test(args, 116 | re.escape(stdout_expected) + r'\Z', 117 | re.escape(stderr_expected) + r'\Z', 118 | 0) 119 | 120 | def main_test(self, args, stdout_pattern, stderr_pattern, exit_code): 121 | # type: (List[str], str, str, int) -> None 122 | save_stdout = sys.stdout 123 | save_stderr = sys.stderr 124 | stdout = StringIO() 125 | stderr = StringIO() 126 | try: 127 | sys.stdout = stdout 128 | sys.stderr = stderr 129 | dunder_main(args) 130 | code = 0 131 | except SystemExit as err: 132 | code = err.code 133 | finally: 134 | sys.stdout = save_stdout 135 | sys.stderr = save_stderr 136 | stdout_value = stdout.getvalue() 137 | stderr_value = stderr.getvalue() 138 | assert re.match(stdout_pattern, stdout_value) 139 | match = re.match(stderr_pattern, stderr_value) 140 | ## if not match: print("\nNah") 141 | ## else: print("\nYa!") 142 | ## print(stderr_value) 143 | ## import pdb; pdb.set_trace() 144 | assert code == exit_code 145 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/infer_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | from typing import List, Tuple 4 | 5 | from pyannotate_tools.annotations.infer import ( 6 | flatten_types, 7 | infer_annotation, 8 | merge_items, 9 | remove_redundant_items, 10 | ) 11 | from pyannotate_tools.annotations.types import ( 12 | AbstractType, 13 | AnyType, 14 | ARG_POS, 15 | ARG_STAR, 16 | ClassType, 17 | TupleType, 18 | UnionType, 19 | NoReturnType, 20 | ) 21 | 22 | 23 | class TestInfer(unittest.TestCase): 24 | def test_simple(self): 25 | # type: () -> None 26 | self.assert_infer(['(int) -> str'], ([(ClassType('int'), ARG_POS)], 27 | ClassType('str'))) 28 | 29 | def test_infer_union_arg(self): 30 | # type: () -> None 31 | self.assert_infer(['(int) -> None', 32 | '(str) -> None'], 33 | ([(UnionType([ClassType('int'), 34 | ClassType('str')]), ARG_POS)], 35 | ClassType('None'))) 36 | 37 | def test_infer_union_return(self): 38 | # type: () -> None 39 | self.assert_infer(['() -> int', 40 | '() -> str'], 41 | ([], 42 | UnionType([ClassType('int'), ClassType('str')]))) 43 | 44 | def test_star_arg(self): 45 | # type: () -> None 46 | self.assert_infer(['(int) -> None', 47 | '(int, *bool) -> None'], 48 | ([(ClassType('int'), ARG_POS), 49 | (ClassType('bool'), ARG_STAR)], 50 | ClassType('None'))) 51 | 52 | def test_merge_unions(self): 53 | # type: () -> None 54 | self.assert_infer(['(Union[int, str]) -> None', 55 | '(Union[str, None]) -> None'], 56 | ([(UnionType([ClassType('int'), 57 | ClassType('str'), 58 | ClassType('None')]), ARG_POS)], 59 | ClassType('None'))) 60 | 61 | def test_remove_redundant_union_item(self): 62 | # type: () -> None 63 | self.assert_infer(['(str) -> None', 64 | '(unicode) -> None'], 65 | ([(ClassType('Text'), ARG_POS)], 66 | ClassType('None'))) 67 | 68 | def test_remove_redundant_dict_item(self): 69 | # type: () -> None 70 | self.assert_infer(['(Dict[str, Any]) -> None', 71 | '(Dict[str, str]) -> None'], 72 | ([(ClassType('Dict', [ClassType('str'), AnyType()]), ARG_POS)], 73 | ClassType('None'))) 74 | 75 | def test_remove_redundant_dict_item_when_simplified(self): 76 | # type: () -> None 77 | self.assert_infer(['(Dict[str, Any]) -> None', 78 | '(Dict[str, Union[str, List, Dict, int]]) -> None'], 79 | ([(ClassType('Dict', [ClassType('str'), AnyType()]), ARG_POS)], 80 | ClassType('None'))) 81 | 82 | def test_simplify_list_item_types(self): 83 | # type: () -> None 84 | self.assert_infer(['(List[Union[bool, int]]) -> None'], 85 | ([(ClassType('List', [ClassType('int')]), ARG_POS)], 86 | ClassType('None'))) 87 | 88 | def test_simplify_potential_typed_dict(self): 89 | # type: () -> None 90 | # Fall back to Dict[x, Any] in case of a complex Dict type. 91 | self.assert_infer(['(Dict[str, Union[int, str]]) -> Any'], 92 | ([(ClassType('Dict', [ClassType('str'), AnyType()]), ARG_POS)], 93 | AnyType())) 94 | self.assert_infer(['(Dict[Text, Union[int, str]]) -> Any'], 95 | ([(ClassType('Dict', [ClassType('Text'), AnyType()]), ARG_POS)], 96 | AnyType())) 97 | # Not a potential TypedDict so ordinary simplification applies. 98 | self.assert_infer(['(Dict[str, Union[str, Text]]) -> Any'], 99 | ([(ClassType('Dict', [ClassType('str'), ClassType('Text')]), ARG_POS)], 100 | AnyType())) 101 | self.assert_infer(['(Dict[str, Union[int, None]]) -> Any'], 102 | ([(ClassType('Dict', [ClassType('str'), 103 | UnionType([ClassType('int'), 104 | ClassType('None')])]), ARG_POS)], 105 | AnyType())) 106 | 107 | def test_simplify_multiple_empty_collections(self): 108 | # type: () -> None 109 | self.assert_infer(['() -> Tuple[List, List[x]]', 110 | '() -> Tuple[List, List]'], 111 | ([], 112 | TupleType([ClassType('List'), ClassType('List', [ClassType('x')])]))) 113 | 114 | def assert_infer(self, comments, expected): 115 | # type: (List[str], Tuple[List[Tuple[AbstractType, str]], AbstractType]) -> None 116 | actual = infer_annotation(comments) 117 | assert actual == expected 118 | 119 | def test_infer_ignore_mock(self): 120 | # type: () -> None 121 | self.assert_infer(['(mock.mock.Mock) -> None', 122 | '(str) -> None'], 123 | ([(ClassType('str'), ARG_POS)], 124 | ClassType('None'))) 125 | 126 | def test_infer_ignore_mock_fallback_to_any(self): 127 | # type: () -> None 128 | self.assert_infer(['(mock.mock.Mock) -> str', 129 | '(mock.mock.Mock) -> int'], 130 | ([(AnyType(), ARG_POS)], 131 | UnionType([ClassType('str'), ClassType('int')]))) 132 | 133 | def test_infer_none_argument(self): 134 | # type: () -> None 135 | self.assert_infer(['(None) -> None'], 136 | ([(UnionType([ClassType('None'), AnyType()]), ARG_POS)], 137 | ClassType('None'))) 138 | 139 | CT = ClassType 140 | 141 | 142 | class TestRedundantItems(unittest.TestCase): 143 | def test_cannot_simplify(self): 144 | # type: () -> None 145 | for first, second in ((CT('str'), CT('int')), 146 | (CT('List', [CT('int')]), 147 | CT('List', [CT('str')])), 148 | (CT('List'), 149 | CT('Set', [CT('int')]))): 150 | assert remove_redundant_items([first, second]) == [first, second] 151 | assert remove_redundant_items([second, first]) == [second, first] 152 | 153 | def test_simplify_simple(self): 154 | # type: () -> None 155 | for first, second in ((CT('str'), CT('Text')), 156 | (CT('bool'), CT('int')), 157 | (CT('int'), CT('float'))): 158 | assert remove_redundant_items([first, second]) == [second] 159 | assert remove_redundant_items([second, first]) == [second] 160 | 161 | def test_simplify_multiple(self): 162 | # type: () -> None 163 | assert remove_redundant_items([CT('Text'), CT('str'), CT('bool'), CT('int'), 164 | CT('X')]) == [CT('Text'), CT('int'), CT('X')] 165 | 166 | def test_simplify_generics(self): 167 | # type: () -> None 168 | for first, second in ((CT('List'), CT('List', [CT('Text')])), 169 | (CT('Set'), CT('Set', [CT('Text')])), 170 | (CT('Dict'), CT('Dict', [CT('str'), CT('int')]))): 171 | assert remove_redundant_items([first, second]) == [second] 172 | 173 | 174 | class TestMergeUnionItems(unittest.TestCase): 175 | def test_cannot_merge(self): 176 | # type: () -> None 177 | for first, second in ((CT('str'), CT('Text')), 178 | (CT('List', [CT('int')]), CT('List', [CT('str')]))): 179 | assert merge_items([first, second]) == [first, second] 180 | assert merge_items([second, first]) == [second, first] 181 | assert merge_items([first, second, first]) == [first, second, first] 182 | 183 | def test_merge_union_of_same_length_tuples(self): 184 | # type: () -> None 185 | assert merge_items([TupleType([CT('str')]), 186 | TupleType([CT('int')])]) == [TupleType([UnionType([CT('str'), 187 | CT('int')])])] 188 | assert merge_items([TupleType([CT('str')]), 189 | TupleType([CT('Text')])]) == [TupleType([CT('Text')])] 190 | 191 | def test_merge_tuples_with_different_lengths(self): 192 | # type: () -> None 193 | assert merge_items([ 194 | TupleType([CT('str')]), 195 | TupleType([CT('str'), CT('str')])]) == [CT('Tuple', [CT('str')])] 196 | assert merge_items([ 197 | TupleType([]), 198 | TupleType([CT('str')]), 199 | TupleType([CT('str'), CT('str')])]) == [CT('Tuple', [CT('str')])] 200 | # Don't merge if types aren't identical 201 | assert merge_items([ 202 | TupleType([CT('str')]), 203 | TupleType([CT('str'), CT('int')])]) == [TupleType([CT('str')]), 204 | TupleType([CT('str'), CT('int')])] 205 | 206 | def test_merge_union_containing_no_return(self): 207 | # type: () -> None 208 | assert merge_items([CT('int'), NoReturnType()]) == [CT('int')] 209 | assert merge_items([NoReturnType(), CT('int')]) == [CT('int')] 210 | 211 | 212 | class TestFlattenTypes(unittest.TestCase): 213 | def test_nested_tuples(self): 214 | # type: () -> None 215 | assert flatten_types([UnionType([UnionType([CT('int'), CT('str')]), CT('X')])]) == [ 216 | CT('int'), CT('str'), CT('X')] 217 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/main_test.py: -------------------------------------------------------------------------------- 1 | import contextlib 2 | import os 3 | import tempfile 4 | import textwrap 5 | import unittest 6 | 7 | from typing import Iterator 8 | 9 | from pyannotate_tools.annotations.infer import InferError 10 | from pyannotate_tools.annotations.main import (generate_annotations_json, 11 | generate_annotations_json_string) 12 | 13 | 14 | class TestMain(unittest.TestCase): 15 | def test_generation(self): 16 | # type: () -> None 17 | data = """ 18 | [ 19 | { 20 | "path": "pkg/thing.py", 21 | "line": 422, 22 | "func_name": "my_function", 23 | "type_comments": [ 24 | "(List[int], str) -> None" 25 | ], 26 | "samples": 3 27 | } 28 | ] 29 | """ 30 | with self.temporary_file() as target_path: 31 | with self.temporary_json_file(data) as source_path: 32 | generate_annotations_json(source_path, target_path) 33 | with open(target_path) as target: 34 | actual = target.read() 35 | 36 | actual = actual.replace(' \n', '\n') 37 | expected = textwrap.dedent("""\ 38 | [ 39 | { 40 | "func_name": "my_function", 41 | "line": 422, 42 | "path": "pkg/thing.py", 43 | "samples": 3, 44 | "signature": { 45 | "arg_types": [ 46 | "List[int]", 47 | "str" 48 | ], 49 | "return_type": "None" 50 | } 51 | } 52 | ]""") 53 | assert actual == expected 54 | 55 | def test_ambiguous_kind(self): 56 | # type: () -> None 57 | data = """ 58 | [ 59 | { 60 | "path": "pkg/thing.py", 61 | "line": 422, 62 | "func_name": "my_function", 63 | "type_comments": [ 64 | "(List[int], str) -> None", 65 | "(List[int], *str) -> None" 66 | ], 67 | "samples": 3 68 | } 69 | ] 70 | """ 71 | with self.assertRaises(InferError) as e: 72 | with self.temporary_json_file(data) as source_path: 73 | generate_annotations_json(source_path, '/dummy') 74 | assert str(e.exception) == textwrap.dedent("""\ 75 | Ambiguous argument kinds: 76 | (List[int], str) -> None 77 | (List[int], *str) -> None""") 78 | 79 | def test_generate_to_memory(self): 80 | # type: () -> None 81 | data = """ 82 | [ 83 | { 84 | "path": "pkg/thing.py", 85 | "line": 422, 86 | "func_name": "my_function", 87 | "type_comments": [ 88 | "(List[int], str) -> None" 89 | ], 90 | "samples": 3 91 | } 92 | ] 93 | """ 94 | with self.temporary_json_file(data) as source_path: 95 | output_data = generate_annotations_json_string(source_path) 96 | assert output_data == [ 97 | { 98 | "path": "pkg/thing.py", 99 | "line": 422, 100 | "func_name": "my_function", 101 | "signature": { 102 | "arg_types": [ 103 | "List[int]", 104 | "str" 105 | ], 106 | "return_type": "None" 107 | }, 108 | "samples": 3 109 | } 110 | ] 111 | 112 | with self.temporary_json_file(data) as source_path: 113 | output_data = generate_annotations_json_string(source_path, only_simple=True) 114 | assert output_data == [] 115 | 116 | def test_generate_simple_signatures(self): 117 | # type: () -> None 118 | data = """ 119 | [ 120 | { 121 | "path": "pkg/thing.py", 122 | "line": 422, 123 | "func_name": "complex_function", 124 | "type_comments": [ 125 | "(List[int], str) -> None" 126 | ], 127 | "samples": 3 128 | }, 129 | { 130 | "path": "pkg/thing.py", 131 | "line": 9000, 132 | "func_name": "simple_function", 133 | "type_comments": [ 134 | "(int, str) -> None" 135 | ], 136 | "samples": 3 137 | } 138 | ] 139 | """ 140 | with self.temporary_json_file(data) as source_path: 141 | output_data = generate_annotations_json_string(source_path, only_simple=True) 142 | assert output_data == [ 143 | { 144 | "path": "pkg/thing.py", 145 | "line": 9000, 146 | "func_name": "simple_function", 147 | "signature": { 148 | "arg_types": [ 149 | "int", 150 | "str" 151 | ], 152 | "return_type": "None" 153 | }, 154 | "samples": 3 155 | } 156 | ] 157 | 158 | @contextlib.contextmanager 159 | def temporary_json_file(self, data): 160 | # type: (str) -> Iterator[str] 161 | source = None 162 | try: 163 | with tempfile.NamedTemporaryFile(mode='w', delete=False) as source: 164 | source.write(data) 165 | yield source.name 166 | finally: 167 | if source is not None: 168 | os.remove(source.name) 169 | 170 | @contextlib.contextmanager 171 | def temporary_file(self): 172 | # type: () -> Iterator[str] 173 | target = None 174 | try: 175 | with tempfile.NamedTemporaryFile(mode='w', delete=False) as target: 176 | pass 177 | yield target.name 178 | finally: 179 | if target is not None: 180 | os.remove(target.name) 181 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/parse_test.py: -------------------------------------------------------------------------------- 1 | import os 2 | import tempfile 3 | import unittest 4 | 5 | from typing import List, Optional, Tuple 6 | 7 | from pyannotate_tools.annotations.parse import parse_json, parse_type_comment, ParseError, tokenize 8 | from pyannotate_tools.annotations.types import ( 9 | AbstractType, 10 | AnyType, 11 | ARG_POS, 12 | ARG_STAR, 13 | ARG_STARSTAR, 14 | Argument, 15 | ClassType, 16 | TupleType, 17 | UnionType, 18 | NoReturnType, 19 | ) 20 | 21 | 22 | class TestParseError(unittest.TestCase): 23 | def test_str_conversion(self): 24 | # type: () -> None 25 | assert str(ParseError('(int -> str')) == 'Invalid type comment: (int -> str' 26 | 27 | 28 | class TestParseJson(unittest.TestCase): 29 | def test_parse_json(self): 30 | # type: () -> None 31 | data = """ 32 | [ 33 | { 34 | "path": "pkg/thing.py", 35 | "line": 422, 36 | "func_name": "my_function", 37 | "type_comments": [ 38 | "(int) -> None", 39 | "(str) -> None" 40 | ], 41 | "samples": 3 42 | } 43 | ] 44 | """ 45 | f = None 46 | try: 47 | with tempfile.NamedTemporaryFile(mode='w', delete=False) as f: 48 | f.write(data) 49 | result = parse_json(f.name) 50 | finally: 51 | if f is not None: 52 | os.remove(f.name) 53 | assert len(result) == 1 54 | item = result[0] 55 | assert item.path == 'pkg/thing.py' 56 | assert item.line == 422 57 | assert item.func_name == 'my_function' 58 | assert item.type_comments == ['(int) -> None', 59 | '(str) -> None'] 60 | assert item.samples == 3 61 | 62 | 63 | class TestTokenize(unittest.TestCase): 64 | def test_tokenize(self): 65 | # type: () -> None 66 | self.assert_tokenize( 67 | ' List[int, str] ( )-> *', 68 | 'DottedName(List) [ DottedName(int) , DottedName(str) ] ( ) -> * End()') 69 | 70 | def test_special_cases(self): 71 | # type: () -> None 72 | self.assert_tokenize('dictionary-itemiterator', 73 | 'DottedName(Iterator) End()') 74 | self.assert_tokenize('dictionary-keyiterator', 75 | 'DottedName(Iterator) End()') 76 | self.assert_tokenize('dictionary-valueiterator', 77 | 'DottedName(Iterator) End()') 78 | self.assert_tokenize('foo-bar', 'DottedName(Any) End()') 79 | self.assert_tokenize('pytz.tzfile.Europe/Amsterdam', 80 | 'DottedName(datetime.tzinfo) End()') 81 | 82 | def assert_tokenize(self, s, expected): 83 | # type: (str, str) -> None 84 | tokens = tokenize(s) 85 | actual = ' '.join(str(t) for t in tokens) 86 | assert actual == expected 87 | 88 | 89 | def class_arg(name, args=None): 90 | # type: (str, Optional[List[AbstractType]]) -> Argument 91 | return Argument(ClassType(name, args), ARG_POS) 92 | 93 | 94 | def any_arg(): 95 | # type: () -> Argument 96 | return Argument(AnyType(), ARG_POS) 97 | 98 | 99 | def tuple_arg(items): 100 | # type: (List[AbstractType]) -> Argument 101 | return Argument(TupleType(items), ARG_POS) 102 | 103 | 104 | 105 | class TestParseTypeComment(unittest.TestCase): 106 | def test_empty(self): 107 | # type: () -> None 108 | self.assert_type_comment('() -> None', ([], ClassType('None'))) 109 | 110 | def test_simple_args(self): 111 | # type: () -> None 112 | self.assert_type_comment('(int) -> None', ([class_arg('int')], ClassType('None'))) 113 | self.assert_type_comment('(int, str) -> bool', ([class_arg('int'), 114 | class_arg('str')], ClassType('bool'))) 115 | 116 | def test_generic(self): 117 | # type: () -> None 118 | self.assert_type_comment('(List[int]) -> Dict[str, bool]', 119 | ([class_arg('List', [ClassType('int')])], 120 | ClassType('Dict', [ClassType('str'), ClassType('bool')]))) 121 | 122 | def test_any_and_unknown(self): 123 | # type: () -> None 124 | self.assert_type_comment('(Any) -> pyannotate_runtime.collect_types.UnknownType', 125 | ([any_arg()], AnyType())) 126 | 127 | def test_no_return(self): 128 | # type: () -> None 129 | self.assert_type_comment('() -> pyannotate_runtime.collect_types.NoReturnType', 130 | ([], NoReturnType())) 131 | 132 | def test_tuple(self): 133 | # type: () -> None 134 | self.assert_type_comment('(Tuple[]) -> Any', ([tuple_arg([])], AnyType())) 135 | self.assert_type_comment('(Tuple[int]) -> Any', 136 | ([tuple_arg([ClassType('int')])], AnyType())) 137 | self.assert_type_comment('(Tuple[int, str]) -> Any', 138 | ([tuple_arg([ClassType('int'), 139 | ClassType('str')])], AnyType())) 140 | 141 | def test_union(self): 142 | # type: () -> None 143 | self.assert_type_comment('(Union[int, str]) -> Any', 144 | ([Argument(UnionType([ClassType('int'), 145 | ClassType('str')]), ARG_POS)], AnyType())) 146 | self.assert_type_comment('(Union[int]) -> Any', 147 | ([class_arg('int')], AnyType())) 148 | 149 | def test_optional(self): 150 | # type: () -> None 151 | self.assert_type_comment('(Optional[int]) -> Any', 152 | ([Argument(UnionType([ClassType('int'), 153 | ClassType('None')]), ARG_POS)], AnyType())) 154 | 155 | def test_star_args(self): 156 | # type: () -> None 157 | self.assert_type_comment('(*str) -> Any', 158 | ([Argument(ClassType('str'), ARG_STAR)], AnyType())) 159 | self.assert_type_comment('(int, *str) -> Any', 160 | ([class_arg('int'), Argument(ClassType('str'), ARG_STAR)], 161 | AnyType())) 162 | 163 | def test_star_star_args(self): 164 | # type: () -> None 165 | self.assert_type_comment('(**str) -> Any', 166 | ([Argument(ClassType('str'), ARG_STARSTAR)], AnyType())) 167 | self.assert_type_comment('(int, *str, **bool) -> Any', 168 | ([class_arg('int'), 169 | Argument(ClassType('str'), ARG_STAR), 170 | Argument(ClassType('bool'), ARG_STARSTAR)], AnyType())) 171 | 172 | def test_function(self): 173 | # type: () -> None 174 | self.assert_type_comment('(function) -> Any', 175 | ([class_arg('Callable')], AnyType())) 176 | 177 | def test_unicode(self): 178 | # type: () -> None 179 | self.assert_type_comment('(unicode) -> Any', 180 | ([class_arg('Text')], AnyType())) 181 | 182 | def test_bad_annotation(self): 183 | # type: () -> None 184 | for bad in ['( -> None', 185 | '()', 186 | ')) -> None', 187 | '() -> ', 188 | '()->', 189 | '() -> None x', 190 | 'int', 191 | 'int -> None', 192 | '(Union[]) -> None', 193 | '(List[int) -> None', 194 | '(*int, *str) -> None', 195 | '(*int, int) -> None', 196 | '(**int, *str) -> None', 197 | '(**int, str) -> None', 198 | '(**int, **str) -> None']: 199 | with self.assertRaises(ParseError): 200 | parse_type_comment(bad) 201 | 202 | def assert_type_comment(self, comment, expected): 203 | # type: (str, Tuple[List[Argument], AbstractType]) -> None 204 | actual = parse_type_comment(comment) 205 | assert actual == expected 206 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/tests/types_test.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | from pyannotate_tools.annotations.types import AnyType, ClassType, TupleType, UnionType 4 | 5 | 6 | class TestTypes(unittest.TestCase): 7 | def test_instance_str(self): 8 | # type: () -> None 9 | assert str(ClassType('int')) == 'int' 10 | assert str(ClassType('List', [ClassType('int')])) == 'List[int]' 11 | assert str(ClassType('Dict', [ClassType('int'), 12 | ClassType('str')])) == 'Dict[int, str]' 13 | 14 | def test_any_type_str(self): 15 | # type: () -> None 16 | assert str(AnyType()) == 'Any' 17 | 18 | def test_tuple_type_str(self): 19 | # type: () -> None 20 | assert str(TupleType([ClassType('int')])) == 'Tuple[int]' 21 | assert str(TupleType([ClassType('int'), 22 | ClassType('str')])) == 'Tuple[int, str]' 23 | assert str(TupleType([])) == 'Tuple[()]' 24 | 25 | def test_union_type_str(Self): 26 | # type: () -> None 27 | assert str(UnionType([ClassType('int'), ClassType('str')])) == 'Union[int, str]' 28 | 29 | def test_optional(Self): 30 | # type: () -> None 31 | assert str(UnionType([ClassType('str'), ClassType('None')])) == 'Optional[str]' 32 | assert str(UnionType([ClassType('None'), ClassType('str')])) == 'Optional[str]' 33 | assert str(UnionType([ClassType('None'), ClassType('str'), 34 | ClassType('int')])) == 'Union[None, str, int]' 35 | 36 | def test_uniform_tuple_str(self): 37 | # type: () -> None 38 | assert str(ClassType('Tuple', [ClassType('int')])) == 'Tuple[int, ...]' 39 | -------------------------------------------------------------------------------- /pyannotate_tools/annotations/types.py: -------------------------------------------------------------------------------- 1 | """Internal representation of type objects.""" 2 | 3 | from typing import NamedTuple, Optional, Sequence 4 | 5 | 6 | class AbstractType(object): 7 | """Abstract base class for types.""" 8 | 9 | 10 | class ClassType(AbstractType): 11 | """A class type, potentially generic (int, List[str], None, ...)""" 12 | 13 | def __init__(self, name, args=None): 14 | # type: (str, Optional[Sequence[AbstractType]]) -> None 15 | self.name = name 16 | if args: 17 | self.args = tuple(args) 18 | else: 19 | self.args = () 20 | 21 | def __repr__(self): 22 | # type: () -> str 23 | if self.name == 'Tuple' and len(self.args) == 1: 24 | return 'Tuple[%s, ...]' % self.args[0] 25 | elif self.args: 26 | return '%s[%s]' % (self.name, ', '.join(str(arg) for arg in self.args)) 27 | else: 28 | return self.name 29 | 30 | def __eq__(self, other): 31 | # type: (object) -> bool 32 | return isinstance(other, ClassType) and self.name == other.name and self.args == other.args 33 | 34 | def __hash__(self): 35 | # type: () -> int 36 | return hash((self.name, self.args)) 37 | 38 | 39 | class AnyType(AbstractType): 40 | """The type Any""" 41 | 42 | def __repr__(self): 43 | # type: () -> str 44 | return 'Any' 45 | 46 | def __eq__(self, other): 47 | # type: (object) -> bool 48 | return isinstance(other, AnyType) 49 | 50 | def __hash__(self): 51 | # type: () -> int 52 | return hash('Any') 53 | 54 | 55 | class NoReturnType(AbstractType): 56 | """The type mypy_extensions.NoReturn""" 57 | 58 | def __repr__(self): 59 | # type: () -> str 60 | return 'mypy_extensions.NoReturn' 61 | 62 | def __eq__(self, other): 63 | # type: (object) -> bool 64 | return isinstance(other, NoReturnType) 65 | 66 | def __hash__(self): 67 | # type: () -> int 68 | return hash('NoReturn') 69 | 70 | 71 | class TupleType(AbstractType): 72 | """Fixed-length tuple Tuple[x, ..., y]""" 73 | 74 | def __init__(self, items): 75 | # type: (Sequence[AbstractType]) -> None 76 | self.items = tuple(items) 77 | 78 | def __repr__(self): 79 | # type: () -> str 80 | if not self.items: 81 | return 'Tuple[()]' # Special case 82 | return 'Tuple[%s]' % ', '.join(str(item) for item in self.items) 83 | 84 | def __eq__(self, other): 85 | # type: (object) -> bool 86 | return isinstance(other, TupleType) and self.items == other.items 87 | 88 | def __hash__(self): 89 | # type: () -> int 90 | return hash(('tuple', self.items)) 91 | 92 | 93 | class UnionType(AbstractType): 94 | """Union[x, ..., y]""" 95 | 96 | def __init__(self, items): 97 | # type: (Sequence[AbstractType]) -> None 98 | self.items = tuple(items) 99 | 100 | def __repr__(self): 101 | # type: () -> str 102 | items = self.items 103 | if len(items) == 2: 104 | if is_none(items[0]): 105 | return 'Optional[%s]' % items[1] 106 | elif is_none(items[1]): 107 | return 'Optional[%s]' % items[0] 108 | return 'Union[%s]' % ', '.join(str(item) for item in items) 109 | 110 | def __eq__(self, other): 111 | # type: (object) -> bool 112 | return isinstance(other, UnionType) and set(self.items) == set(other.items) 113 | 114 | def __hash__(self): 115 | # type: () -> int 116 | return hash(('union', self.items)) 117 | 118 | 119 | # Argument kind 120 | ARG_POS = 'ARG_POS' # Normal 121 | ARG_STAR = 'ARG_STAR' # *args 122 | ARG_STARSTAR = 'ARG_STARSTAR' # **kwargs 123 | 124 | # Description of an argument in a signature. The kind is one of ARG_*. 125 | Argument = NamedTuple('Argument', [('type', AbstractType), ('kind', str)]) 126 | 127 | 128 | def is_none(t): 129 | # type: (AbstractType) -> bool 130 | return isinstance(t, ClassType) and t.name == 'None' 131 | 132 | 133 | def is_optional(t): 134 | # type: (AbstractType) -> bool 135 | return (isinstance(t, UnionType) 136 | and len(t.items) == 2 137 | and any(item == ClassType('None') for item in t.items)) 138 | -------------------------------------------------------------------------------- /pyannotate_tools/fixes/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dropbox/pyannotate/a7a46f394f0ba91a1b5fbf657e2393af542969ae/pyannotate_tools/fixes/__init__.py -------------------------------------------------------------------------------- /pyannotate_tools/fixes/fix_annotate.py: -------------------------------------------------------------------------------- 1 | """Fixer that inserts mypy annotations into all methods. 2 | 3 | This transforms e.g. 4 | 5 | def foo(self, bar, baz=12): 6 | return bar + baz 7 | 8 | into a type annoted version: 9 | 10 | def foo(self, bar, baz=12): 11 | # type: (Any, int) -> Any # noqa: F821 12 | return bar + baz 13 | 14 | or (when setting options['annotation_style'] to 'py3'): 15 | 16 | def foo(self, bar : Any, baz : int = 12) -> Any: 17 | return bar + baz 18 | 19 | 20 | It does not do type inference but it recognizes some basic default 21 | argument values such as numbers and strings (and assumes their type 22 | implies the argument type). 23 | 24 | It also uses some basic heuristics to decide whether to ignore the 25 | first argument: 26 | 27 | - always if it's named 'self' 28 | - if there's a @classmethod decorator 29 | 30 | Finally, it knows that __init__() is supposed to return None. 31 | """ 32 | 33 | from __future__ import print_function 34 | 35 | import os 36 | import re 37 | 38 | from lib2to3.fixer_base import BaseFix 39 | from lib2to3.fixer_util import syms, touch_import, find_indentation 40 | from lib2to3.patcomp import compile_pattern 41 | from lib2to3.pgen2 import token 42 | from lib2to3.pytree import Leaf, Node 43 | 44 | 45 | class FixAnnotate(BaseFix): 46 | 47 | # This fixer is compatible with the bottom matcher. 48 | BM_compatible = True 49 | 50 | # This fixer shouldn't run by default. 51 | explicit = True 52 | 53 | # The pattern to match. 54 | PATTERN = """ 55 | funcdef< 'def' name=any parameters=parameters< '(' [args=any] rpar=')' > ':' suite=any+ > 56 | """ 57 | 58 | _maxfixes = os.getenv('MAXFIXES') 59 | counter = None if not _maxfixes else int(_maxfixes) 60 | 61 | def transform(self, node, results): 62 | if FixAnnotate.counter is not None: 63 | if FixAnnotate.counter <= 0: 64 | return 65 | 66 | # Check if there's already a long-form annotation for some argument. 67 | parameters = results.get('parameters') 68 | if parameters is not None: 69 | for ch in parameters.pre_order(): 70 | if ch.prefix.lstrip().startswith('# type:'): 71 | return 72 | args = results.get('args') 73 | if args is not None: 74 | for ch in args.pre_order(): 75 | if ch.prefix.lstrip().startswith('# type:'): 76 | return 77 | 78 | children = results['suite'][0].children 79 | 80 | # NOTE: I've reverse-engineered the structure of the parse tree. 81 | # It's always a list of nodes, the first of which contains the 82 | # entire suite. Its children seem to be: 83 | # 84 | # [0] NEWLINE 85 | # [1] INDENT 86 | # [2...n-2] statements (the first may be a docstring) 87 | # [n-1] DEDENT 88 | # 89 | # Comments before the suite are part of the INDENT's prefix. 90 | # 91 | # "Compact" functions (e.g. "def foo(x, y): return max(x, y)") 92 | # have a different structure (no NEWLINE, INDENT, or DEDENT). 93 | 94 | # Check if there's already an annotation. 95 | for ch in children: 96 | if ch.prefix.lstrip().startswith('# type:'): 97 | return # There's already a # type: comment here; don't change anything. 98 | 99 | # Python 3 style return annotation are already skipped by the pattern 100 | 101 | ### Python 3 style argument annotation structure 102 | # 103 | # Structure of the arguments tokens for one positional argument without default value : 104 | # + LPAR '(' 105 | # + NAME_NODE_OR_LEAF arg1 106 | # + RPAR ')' 107 | # 108 | # NAME_NODE_OR_LEAF is either: 109 | # 1. Just a leaf with value NAME 110 | # 2. A node with children: NAME, ':", node expr or value leaf 111 | # 112 | # Structure of the arguments tokens for one args with default value or multiple 113 | # args, with or without default value, and/or with extra arguments : 114 | # + LPAR '(' 115 | # + node 116 | # [ 117 | # + NAME_NODE_OR_LEAF 118 | # [ 119 | # + EQUAL '=' 120 | # + node expr or value leaf 121 | # ] 122 | # ( 123 | # + COMMA ',' 124 | # + NAME_NODE_OR_LEAF positional argn 125 | # [ 126 | # + EQUAL '=' 127 | # + node expr or value leaf 128 | # ] 129 | # )* 130 | # ] 131 | # [ 132 | # + STAR '*' 133 | # [ 134 | # + NAME_NODE_OR_LEAF positional star argument name 135 | # ] 136 | # ] 137 | # [ 138 | # + COMMA ',' 139 | # + DOUBLESTAR '**' 140 | # + NAME_NODE_OR_LEAF positional keyword argument name 141 | # ] 142 | # + RPAR ')' 143 | 144 | # Let's skip Python 3 argument annotations 145 | it = iter(args.children) if args else iter([]) 146 | for ch in it: 147 | if ch.type == token.STAR: 148 | # *arg part 149 | ch = next(it) 150 | if ch.type == token.COMMA: 151 | continue 152 | elif ch.type == token.DOUBLESTAR: 153 | # *arg part 154 | ch = next(it) 155 | if ch.type > 256: 156 | # this is a node, therefore an annotation 157 | assert ch.children[0].type == token.NAME 158 | return 159 | try: 160 | ch = next(it) 161 | if ch.type == token.COLON: 162 | # this is an annotation 163 | return 164 | elif ch.type == token.EQUAL: 165 | ch = next(it) 166 | ch = next(it) 167 | assert ch.type == token.COMMA 168 | continue 169 | except StopIteration: 170 | break 171 | 172 | # Compute the annotation 173 | annot = self.make_annotation(node, results) 174 | if annot is None: 175 | return 176 | argtypes, restype = annot 177 | 178 | if self.options['annotation_style'] == 'py3': 179 | self.add_py3_annot(argtypes, restype, node, results) 180 | else: 181 | self.add_py2_annot(argtypes, restype, node, results) 182 | 183 | # Common to py2 and py3 style annotations: 184 | if FixAnnotate.counter is not None: 185 | FixAnnotate.counter -= 1 186 | 187 | # Also add 'from typing import Any' at the top if needed. 188 | self.patch_imports(argtypes + [restype], node) 189 | 190 | def add_py3_annot(self, argtypes, restype, node, results): 191 | args = results.get('args') 192 | 193 | argleaves = [] 194 | if args is None: 195 | # function with 0 arguments 196 | it = iter([]) 197 | elif len(args.children) == 0: 198 | # function with 1 argument 199 | it = iter([args]) 200 | else: 201 | # function with multiple arguments or 1 arg with default value 202 | it = iter(args.children) 203 | 204 | for ch in it: 205 | argstyle = 'name' 206 | if ch.type == token.STAR: 207 | # *arg part 208 | argstyle = 'star' 209 | ch = next(it) 210 | if ch.type == token.COMMA: 211 | continue 212 | elif ch.type == token.DOUBLESTAR: 213 | # *arg part 214 | argstyle = 'keyword' 215 | ch = next(it) 216 | assert ch.type == token.NAME 217 | argleaves.append((argstyle, ch)) 218 | try: 219 | ch = next(it) 220 | if ch.type == token.EQUAL: 221 | ch = next(it) 222 | ch = next(it) 223 | assert ch.type == token.COMMA 224 | continue 225 | except StopIteration: 226 | break 227 | 228 | # when self or cls is not annotated, argleaves == argtypes+1 229 | argleaves = argleaves[len(argleaves) - len(argtypes):] 230 | 231 | for ch_withstyle, chtype in zip(argleaves, argtypes): 232 | style, ch = ch_withstyle 233 | if style == 'star': 234 | assert chtype[0] == '*' 235 | assert chtype[1] != '*' 236 | chtype = chtype[1:] 237 | elif style == 'keyword': 238 | assert chtype[0:2] == '**' 239 | assert chtype[2] != '*' 240 | chtype = chtype[2:] 241 | ch.value = '%s: %s' % (ch.value, chtype) 242 | 243 | # put spaces around the equal sign 244 | if ch.next_sibling and ch.next_sibling.type == token.EQUAL: 245 | nextch = ch.next_sibling 246 | if not nextch.prefix[:1].isspace(): 247 | nextch.prefix = ' ' + nextch.prefix 248 | nextch = nextch.next_sibling 249 | assert nextch != None 250 | if not nextch.prefix[:1].isspace(): 251 | nextch.prefix = ' ' + nextch.prefix 252 | 253 | # Add return annotation 254 | rpar = results['rpar'] 255 | rpar.value = '%s -> %s' % (rpar.value, restype) 256 | 257 | rpar.changed() 258 | 259 | def add_py2_annot(self, argtypes, restype, node, results): 260 | children = results['suite'][0].children 261 | 262 | # Insert '# type: {annot}' comment. 263 | # For reference, see lib2to3/fixes/fix_tuple_params.py in stdlib. 264 | if len(children) >= 1 and children[0].type != token.NEWLINE: 265 | # one liner function 266 | if children[0].prefix.strip() == '': 267 | children[0].prefix = '' 268 | children.insert(0, Leaf(token.NEWLINE, '\n')) 269 | children.insert( 270 | 1, Leaf(token.INDENT, find_indentation(node) + ' ')) 271 | children.append(Leaf(token.DEDENT, '')) 272 | if len(children) >= 2 and children[1].type == token.INDENT: 273 | degen_str = '(...) -> %s' % restype 274 | short_str = '(%s) -> %s' % (', '.join(argtypes), restype) 275 | if (len(short_str) > 64 or len(argtypes) > 5) and len(short_str) > len(degen_str): 276 | self.insert_long_form(node, results, argtypes) 277 | annot_str = degen_str 278 | else: 279 | annot_str = short_str 280 | children[1].prefix = '%s# type: %s\n%s' % (children[1].value, annot_str, 281 | children[1].prefix) 282 | children[1].changed() 283 | else: 284 | self.log_message("%s:%d: cannot insert annotation for one-line function" % 285 | (self.filename, node.get_lineno())) 286 | 287 | def insert_long_form(self, node, results, argtypes): 288 | argtypes = list(argtypes) # We destroy it 289 | args = results['args'] 290 | if isinstance(args, Node): 291 | children = args.children 292 | elif isinstance(args, Leaf): 293 | children = [args] 294 | else: 295 | children = [] 296 | # Interpret children according to the following grammar: 297 | # (('*'|'**')? NAME ['=' expr] ','?)* 298 | flag = False # Set when the next leaf should get a type prefix 299 | indent = '' # Will be set by the first child 300 | 301 | def set_prefix(child): 302 | if argtypes: 303 | arg = argtypes.pop(0).lstrip('*') 304 | else: 305 | arg = 'Any' # Somehow there aren't enough args 306 | if not arg: 307 | # Skip self (look for 'check_self' below) 308 | prefix = child.prefix.rstrip() 309 | else: 310 | prefix = ' # type: ' + arg 311 | old_prefix = child.prefix.strip() 312 | if old_prefix: 313 | assert old_prefix.startswith('#') 314 | prefix += ' ' + old_prefix 315 | child.prefix = prefix + '\n' + indent 316 | 317 | check_self = self.is_method(node) 318 | for child in children: 319 | if check_self and isinstance(child, Leaf) and child.type == token.NAME: 320 | check_self = False 321 | if child.value in ('self', 'cls'): 322 | argtypes.insert(0, '') 323 | if not indent: 324 | indent = ' ' * child.column 325 | if isinstance(child, Leaf) and child.value == ',': 326 | flag = True 327 | elif isinstance(child, Leaf) and flag: 328 | set_prefix(child) 329 | flag = False 330 | need_comma = len(children) >= 1 and children[-1].type != token.COMMA 331 | if need_comma and len(children) >= 2: 332 | if (children[-1].type == token.NAME and 333 | (children[-2].type in (token.STAR, token.DOUBLESTAR))): 334 | need_comma = False 335 | if need_comma: 336 | children.append(Leaf(token.COMMA, u",")) 337 | # Find the ')' and insert a prefix before it too. 338 | parameters = args.parent 339 | close_paren = parameters.children[-1] 340 | assert close_paren.type == token.RPAR, close_paren 341 | set_prefix(close_paren) 342 | assert not argtypes, argtypes 343 | 344 | def patch_imports(self, types, node): 345 | for typ in types: 346 | if 'Any' in typ: 347 | touch_import('typing', 'Any', node) 348 | break 349 | 350 | def make_annotation(self, node, results): 351 | name = results['name'] 352 | assert isinstance(name, Leaf), repr(name) 353 | assert name.type == token.NAME, repr(name) 354 | decorators = self.get_decorators(node) 355 | is_method = self.is_method(node) 356 | if name.value == '__init__' or not self.has_return_exprs(node): 357 | restype = 'None' 358 | else: 359 | restype = 'Any' 360 | args = results.get('args') 361 | argtypes = [] 362 | if isinstance(args, Node): 363 | children = args.children 364 | elif isinstance(args, Leaf): 365 | children = [args] 366 | else: 367 | children = [] 368 | # Interpret children according to the following grammar: 369 | # (('*'|'**')? NAME ['=' expr] ','?)* 370 | stars = inferred_type = '' 371 | in_default = False 372 | at_start = True 373 | for child in children: 374 | if isinstance(child, Leaf): 375 | if child.value in ('*', '**'): 376 | stars += child.value 377 | elif child.type == token.NAME and not in_default: 378 | if not is_method or not at_start or 'staticmethod' in decorators: 379 | inferred_type = 'Any' 380 | else: 381 | # Always skip the first argument if it's named 'self'. 382 | # Always skip the first argument of a class method. 383 | if child.value == 'self' or 'classmethod' in decorators: 384 | pass 385 | else: 386 | inferred_type = 'Any' 387 | elif child.value == '=': 388 | in_default = True 389 | elif in_default and child.value != ',': 390 | if child.type == token.NUMBER: 391 | if re.match(r'\d+[lL]?$', child.value): 392 | inferred_type = 'int' 393 | else: 394 | inferred_type = 'float' # TODO: complex? 395 | elif child.type == token.STRING: 396 | if child.value.startswith(('u', 'U')): 397 | inferred_type = 'unicode' 398 | else: 399 | inferred_type = 'str' 400 | elif child.type == token.NAME and child.value in ('True', 'False'): 401 | inferred_type = 'bool' 402 | elif child.value == ',': 403 | if inferred_type: 404 | argtypes.append(stars + inferred_type) 405 | # Reset 406 | stars = inferred_type = '' 407 | in_default = False 408 | at_start = False 409 | if inferred_type: 410 | argtypes.append(stars + inferred_type) 411 | return argtypes, restype 412 | 413 | # The parse tree has a different shape when there is a single 414 | # decorator vs. when there are multiple decorators. 415 | DECORATED = "decorated< (d=decorator | decorators< dd=decorator+ >) funcdef >" 416 | decorated = compile_pattern(DECORATED) 417 | 418 | def get_decorators(self, node): 419 | """Return a list of decorators found on a function definition. 420 | 421 | This is a list of strings; only simple decorators 422 | (e.g. @staticmethod) are returned. 423 | 424 | If the function is undecorated or only non-simple decorators 425 | are found, return []. 426 | """ 427 | if node.parent is None: 428 | return [] 429 | results = {} 430 | if not self.decorated.match(node.parent, results): 431 | return [] 432 | decorators = results.get('dd') or [results['d']] 433 | decs = [] 434 | for d in decorators: 435 | for child in d.children: 436 | if isinstance(child, Leaf) and child.type == token.NAME: 437 | decs.append(child.value) 438 | return decs 439 | 440 | def is_method(self, node): 441 | """Return whether the node occurs (directly) inside a class.""" 442 | node = node.parent 443 | while node is not None: 444 | if node.type == syms.classdef: 445 | return True 446 | if node.type == syms.funcdef: 447 | return False 448 | node = node.parent 449 | return False 450 | 451 | RETURN_EXPR = "return_stmt< 'return' any >" 452 | return_expr = compile_pattern(RETURN_EXPR) 453 | 454 | def has_return_exprs(self, node): 455 | """Traverse the tree below node looking for 'return expr'. 456 | 457 | Return True if at least 'return expr' is found, False if not. 458 | (If both 'return' and 'return expr' are found, return True.) 459 | """ 460 | results = {} 461 | if self.return_expr.match(node, results): 462 | return True 463 | for child in node.children: 464 | if child.type not in (syms.funcdef, syms.classdef): 465 | if self.has_return_exprs(child): 466 | return True 467 | return False 468 | 469 | YIELD_EXPR = "yield_expr< 'yield' [any] >" 470 | yield_expr = compile_pattern(YIELD_EXPR) 471 | 472 | def is_generator(self, node): 473 | """Traverse the tree below node looking for 'yield [expr]'.""" 474 | results = {} 475 | if self.yield_expr.match(node, results): 476 | return True 477 | for child in node.children: 478 | if child.type not in (syms.funcdef, syms.classdef): 479 | if self.is_generator(child): 480 | return True 481 | return False 482 | -------------------------------------------------------------------------------- /pyannotate_tools/fixes/fix_annotate_json.py: -------------------------------------------------------------------------------- 1 | """Fixer that inserts mypy annotations from json file into code. 2 | 3 | This fixer consumes json from TYPE_COLLECTION_JSON env variable in the following format: 4 | 5 | [ 6 | { 7 | "path": "/Users/svorobev/src/client/build_number/__init__.py", 8 | "func_name": "is_test", 9 | "arg_types": ["int", "str"], 10 | "ret_type": "Any" 11 | }, 12 | ... 13 | ] 14 | 15 | (The old format with "type_comment" instead of "arg_types" and 16 | "ret_type" is also still supported.) 17 | """ 18 | 19 | from __future__ import print_function 20 | 21 | import json # noqa 22 | import os 23 | import re 24 | from contextlib import contextmanager 25 | 26 | from lib2to3.fixer_util import syms, touch_import 27 | from lib2to3.pgen2 import token 28 | from lib2to3.pytree import Base, Leaf, Node 29 | from typing import __all__ as typing_all # type: ignore 30 | from typing import Any, Dict, List, Optional, Tuple 31 | try: 32 | from typing import Text 33 | except ImportError: 34 | # In Python 3.5.1 stdlib, typing.py does not define Text 35 | Text = str # type: ignore 36 | 37 | from .fix_annotate import FixAnnotate 38 | 39 | # Taken from mypy codebase: 40 | # https://github.com/python/mypy/blob/745d300b8304c3dcf601477762bf9d70b9a4619c/mypy/main.py#L503 41 | 42 | PY_EXTENSIONS = ['.pyi', '.py'] 43 | 44 | def crawl_up(arg): 45 | # type: (str) -> Tuple[str, str] 46 | """Given a .py[i] filename, return (root directory, module). 47 | We crawl up the path until we find a directory without 48 | __init__.py[i], or until we run out of path components. 49 | """ 50 | dir, mod = os.path.split(arg) 51 | mod = strip_py(mod) or mod 52 | while dir and get_init_file(dir): 53 | dir, base = os.path.split(dir) 54 | if not base: 55 | break 56 | if mod == '__init__' or not mod: 57 | mod = base 58 | else: 59 | mod = base + '.' + mod 60 | return dir, mod 61 | 62 | def strip_py(arg): 63 | # type: (str) -> Optional[str] 64 | """Strip a trailing .py or .pyi suffix. 65 | Return None if no such suffix is found. 66 | """ 67 | for ext in PY_EXTENSIONS: 68 | if arg.endswith(ext): 69 | return arg[:-len(ext)] 70 | return None 71 | 72 | def get_init_file(dir): 73 | # type: (str) -> Optional[str] 74 | """Check whether a directory contains a file named __init__.py[i]. 75 | If so, return the file's name (with dir prefixed). If not, return 76 | None. 77 | This prefers .pyi over .py (because of the ordering of PY_EXTENSIONS). 78 | """ 79 | for ext in PY_EXTENSIONS: 80 | f = os.path.join(dir, '__init__' + ext) 81 | if os.path.isfile(f): 82 | return f 83 | return None 84 | 85 | def get_funcname(node): 86 | # type: (Optional[Node]) -> Text 87 | """Get function name by (approximately) the following rules: 88 | 89 | - function -> function_name 90 | - method -> ClassName.function_name 91 | 92 | More specifically, we include every class and function name that 93 | the node is a child of, so nested classes and functions get names like 94 | OuterClass.InnerClass.outer_fn.inner_fn. 95 | """ 96 | components = [] # type: List[str] 97 | while node: 98 | if node.type in (syms.classdef, syms.funcdef): 99 | name = node.children[1] 100 | assert name.type == token.NAME, repr(name) 101 | assert isinstance(name, Leaf) # Same as previous, for mypy 102 | components.append(name.value) 103 | node = node.parent 104 | return '.'.join(reversed(components)) 105 | 106 | def count_args(node, results): 107 | # type: (Node, Dict[str, Base]) -> Tuple[int, bool, bool, bool] 108 | """Count arguments and check for self and *args, **kwds. 109 | 110 | Return (selfish, count, star, starstar) where: 111 | - count is total number of args (including *args, **kwds) 112 | - selfish is True if the initial arg is named 'self' or 'cls' 113 | - star is True iff *args is found 114 | - starstar is True iff **kwds is found 115 | """ 116 | count = 0 117 | selfish = False 118 | star = False 119 | starstar = False 120 | args = results.get('args') 121 | if isinstance(args, Node): 122 | children = args.children 123 | elif isinstance(args, Leaf): 124 | children = [args] 125 | else: 126 | children = [] 127 | # Interpret children according to the following grammar: 128 | # (('*'|'**')? NAME ['=' expr] ','?)* 129 | skip = False 130 | previous_token_is_star = False 131 | for child in children: 132 | if skip: 133 | skip = False 134 | elif isinstance(child, Leaf): 135 | # A single '*' indicates the rest of the arguments are keyword only 136 | # and shouldn't be counted as a `*`. 137 | if child.type == token.STAR: 138 | previous_token_is_star = True 139 | elif child.type == token.DOUBLESTAR: 140 | starstar = True 141 | elif child.type == token.NAME: 142 | if count == 0: 143 | if child.value in ('self', 'cls'): 144 | selfish = True 145 | count += 1 146 | if previous_token_is_star: 147 | star = True 148 | elif child.type == token.EQUAL: 149 | skip = True 150 | if child.type != token.STAR: 151 | previous_token_is_star = False 152 | return count, selfish, star, starstar 153 | 154 | class FixAnnotateJson(FixAnnotate): 155 | 156 | needed_imports = None 157 | line_drift = 5 158 | 159 | def add_import(self, mod, name): 160 | if mod == self.current_module(): 161 | return 162 | if self.needed_imports is None: 163 | self.needed_imports = set() 164 | self.needed_imports.add((mod, name)) 165 | 166 | def patch_imports(self, types, node): 167 | if self.needed_imports: 168 | for mod, name in sorted(self.needed_imports): 169 | touch_import(mod, name, node) 170 | self.needed_imports = None 171 | 172 | def set_filename(self, filename): 173 | super(FixAnnotateJson, self).set_filename(filename) 174 | self._current_module = crawl_up(filename)[1] 175 | 176 | def current_module(self): 177 | return self._current_module 178 | 179 | def make_annotation(self, node, results): 180 | name = results['name'] 181 | assert isinstance(name, Leaf), repr(name) 182 | assert name.type == token.NAME, repr(name) 183 | funcname = get_funcname(node) 184 | res = self.get_annotation_from_stub(node, results, funcname) 185 | 186 | # If we couldn't find an annotation and this is a classmethod or 187 | # staticmethod, try again with just the funcname, since the 188 | # type collector can't figure out class names for those. 189 | # (We try with the full name above first so that tools that *can* figure 190 | # that out, like dmypy suggest, can use it.) 191 | if not res: 192 | decs = self.get_decorators(node) 193 | if 'staticmethod' in decs or 'classmethod' in decs: 194 | res = self.get_annotation_from_stub(node, results, name.value) 195 | 196 | return res 197 | 198 | stub_json_file = os.getenv('TYPE_COLLECTION_JSON') 199 | # JSON data for the current file 200 | stub_json = None # type: List[Dict[str, Any]] 201 | 202 | @classmethod 203 | @contextmanager 204 | def max_line_drift_set(cls, max_drift): 205 | old_drift = cls.line_drift 206 | cls.line_drift = max_drift 207 | try: 208 | yield 209 | finally: 210 | cls.line_drift = old_drift 211 | 212 | @classmethod 213 | def init_stub_json_from_data(cls, data, filename): 214 | cls.stub_json = data 215 | cls.top_dir = crawl_up(os.path.abspath(filename))[0] 216 | 217 | def init_stub_json(self): 218 | with open(self.__class__.stub_json_file) as f: 219 | data = json.load(f) 220 | self.__class__.init_stub_json_from_data(data, self.filename) 221 | 222 | def get_annotation_from_stub(self, node, results, funcname): 223 | if not self.__class__.stub_json: 224 | self.init_stub_json() 225 | data = self.__class__.stub_json 226 | # We are using relative paths in the JSON. 227 | items = [ 228 | it for it in data 229 | if it['func_name'] == funcname and 230 | (it['path'] == self.filename or 231 | os.path.join(self.__class__.top_dir, it['path']) == os.path.abspath(self.filename)) 232 | ] 233 | if len(items) > 1: 234 | # this can happen, because of 235 | # 1) nested functions 236 | # 2) method decorators 237 | # as a cheap and dirty solution we just return the nearest one by the line number 238 | # (keep the commented-out log_message call in case we need to come back to this) 239 | ## self.log_message("%s:%d: duplicate signatures for %s (at lines %s)" % 240 | ## (items[0]['path'], node.get_lineno(), items[0]['func_name'], 241 | ## ", ".join(str(it['line']) for it in items))) 242 | items.sort(key=lambda it: abs(node.get_lineno() - it['line'])) 243 | if items: 244 | it = items[0] 245 | # If the line number is too far off, the source probably drifted 246 | # since the trace was collected; it's better to skip this node. 247 | # (Allow some drift, since decorators also cause an offset.) 248 | if abs(node.get_lineno() - it['line']) >= self.line_drift: 249 | self.log_message("%s:%d: '%s' signature from line %d too far away -- skipping" % 250 | (self.filename, node.get_lineno(), it['func_name'], it['line'])) 251 | return None 252 | if 'signature' in it: 253 | sig = it['signature'] 254 | arg_types = sig['arg_types'] 255 | # Passes 1-2 don't always understand *args or **kwds, 256 | # so add '*Any' or '**Any' at the end if needed. 257 | count, selfish, star, starstar = count_args(node, results) 258 | for arg_type in arg_types: 259 | if arg_type.startswith('**'): 260 | starstar = False 261 | elif arg_type.startswith('*'): 262 | star = False 263 | if star: 264 | arg_types.append('*Any') 265 | if starstar: 266 | arg_types.append('**Any') 267 | # Pass 1 omits the first arg iff it's named 'self' or 'cls', 268 | # even if it's not a method, so insert `Any` as needed 269 | # (but only if it's not actually a method). 270 | if selfish and len(arg_types) == count - 1: 271 | if self.is_method(node): 272 | count -= 1 # Leave out the type for 'self' or 'cls' 273 | else: 274 | arg_types.insert(0, 'Any') 275 | # If after those adjustments the count is still off, 276 | # print a warning and skip this node. 277 | if len(arg_types) != count: 278 | self.log_message("%s:%d: source has %d args, annotation has %d -- skipping" % 279 | (self.filename, node.get_lineno(), count, len(arg_types))) 280 | return None 281 | ret_type = sig['return_type'] 282 | arg_types = [self.update_type_names(arg_type) for arg_type in arg_types] 283 | # Avoid common error "No return value expected" 284 | if ret_type == 'None' and self.has_return_exprs(node): 285 | ret_type = 'Optional[Any]' 286 | # Special case for generators. 287 | if (self.is_generator(node) and 288 | not (ret_type == 'Iterator' or ret_type.startswith('Iterator['))): 289 | if ret_type.startswith('Optional['): 290 | assert ret_type[-1] == ']' 291 | ret_type = ret_type[9:-1] 292 | ret_type = 'Iterator[%s]' % ret_type 293 | ret_type = self.update_type_names(ret_type) 294 | return arg_types, ret_type 295 | return None 296 | 297 | def update_type_names(self, type_str): 298 | # Replace e.g. `List[pkg.mod.SomeClass]` with 299 | # `List[SomeClass]` and remember to import it. 300 | return re.sub(r'[\w.:]+', self.type_updater, type_str) 301 | 302 | def type_updater(self, match): 303 | # Replace `pkg.mod.SomeClass` with `SomeClass` 304 | # and remember to import it. 305 | word = match.group() 306 | if word == '...': 307 | return word 308 | if '.' not in word and ':' not in word: 309 | # Assume it's either builtin or from `typing` 310 | if word in typing_all: 311 | self.add_import('typing', word) 312 | return word 313 | # If there is a :, treat that as the separator between the 314 | # module and the class. Otherwise assume everything but the 315 | # last element is the module. 316 | if ':' in word: 317 | mod, name = word.split(':') 318 | to_import = name.split('.', 1)[0] 319 | else: 320 | mod, name = word.rsplit('.', 1) 321 | to_import = name 322 | self.add_import(mod, to_import) 323 | return name 324 | -------------------------------------------------------------------------------- /pyannotate_tools/fixes/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/dropbox/pyannotate/a7a46f394f0ba91a1b5fbf657e2393af542969ae/pyannotate_tools/fixes/tests/__init__.py -------------------------------------------------------------------------------- /pyannotate_tools/fixes/tests/test_annotate_json_py3.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | # Our flake extension misfires on type comments in strings below. 3 | 4 | import json 5 | import os 6 | import tempfile 7 | import unittest 8 | import sys 9 | from mock import patch 10 | 11 | from lib2to3.tests.test_fixers import FixerTestCase 12 | 13 | from pyannotate_tools.fixes.fix_annotate_json import FixAnnotateJson 14 | 15 | 16 | class TestFixAnnotateJson(FixerTestCase): 17 | 18 | def setUp(self): 19 | super(TestFixAnnotateJson, self).setUp( 20 | fix_list=["annotate_json"], 21 | fixer_pkg="pyannotate_tools", 22 | options={'annotation_style' : 'py3'}, 23 | ) 24 | # See https://bugs.python.org/issue14243 for details 25 | self.tf = tempfile.NamedTemporaryFile(mode='w', delete=False) 26 | FixAnnotateJson.stub_json_file = self.tf.name 27 | FixAnnotateJson.stub_json = None 28 | 29 | def tearDown(self): 30 | FixAnnotateJson.stub_json = None 31 | FixAnnotateJson.stub_json_file = None 32 | self.tf.close() 33 | os.remove(self.tf.name) 34 | super(TestFixAnnotateJson, self).tearDown() 35 | 36 | def setTestData(self, data): 37 | json.dump(data, self.tf) 38 | self.tf.close() 39 | self.filename = data[0]["path"] 40 | 41 | def test_basic(self): 42 | self.setTestData( 43 | [{"func_name": "nop", 44 | "path": "", 45 | "line": 3, 46 | "signature": { 47 | "arg_types": ["Foo", "Bar"], 48 | "return_type": "Any"}, 49 | }]) 50 | a = """\ 51 | class Foo: pass 52 | class Bar: pass 53 | def nop(foo, bar): 54 | return 42 55 | """ 56 | b = """\ 57 | from typing import Any 58 | class Foo: pass 59 | class Bar: pass 60 | def nop(foo: Foo, bar: Bar) -> Any: 61 | return 42 62 | """ 63 | self.check(a, b) 64 | 65 | def test_keyword_only_argument(self): 66 | self.setTestData( 67 | [{"func_name": "nop", 68 | "path": "", 69 | "line": 3, 70 | "signature": { 71 | "arg_types": ["Foo", "Bar"], 72 | "return_type": "Any"}, 73 | }]) 74 | a = """\ 75 | class Foo: pass 76 | class Bar: pass 77 | def nop(foo, *, bar): 78 | return 42 79 | """ 80 | b = """\ 81 | from typing import Any 82 | class Foo: pass 83 | class Bar: pass 84 | def nop(foo: Foo, *, bar: Bar) -> Any: 85 | return 42 86 | """ 87 | self.check(a, b) 88 | 89 | def test_add_typing_import(self): 90 | self.setTestData( 91 | [{"func_name": "nop", 92 | "path": "", 93 | "line": 1, 94 | # Check with and without 'typing.' prefix 95 | "signature": { 96 | "arg_types": ["List[typing.AnyStr]", "Callable[[], int]"], 97 | "return_type": "object"}, 98 | }]) 99 | a = """\ 100 | def nop(foo, bar): 101 | return 42 102 | """ 103 | b = """\ 104 | from typing import AnyStr 105 | from typing import Callable 106 | from typing import List 107 | def nop(foo: List[AnyStr], bar: Callable[[], int]) -> object: 108 | return 42 109 | """ 110 | self.check(a, b) 111 | 112 | def test_add_other_import(self): 113 | self.setTestData( 114 | [{"func_name": "nop", 115 | "path": "mod1.py", 116 | "line": 1, 117 | "signature": { 118 | "arg_types": ["mod1.MyClass", "mod2.OtherClass"], 119 | "return_type": "mod3.AnotherClass"}, 120 | }]) 121 | a = """\ 122 | def nop(foo, bar): 123 | return AnotherClass() 124 | class MyClass: pass 125 | """ 126 | b = """\ 127 | from mod2 import OtherClass 128 | from mod3 import AnotherClass 129 | def nop(foo: MyClass, bar: OtherClass) -> AnotherClass: 130 | return AnotherClass() 131 | class MyClass: pass 132 | """ 133 | self.check(a, b) 134 | 135 | def test_add_kwds(self): 136 | self.setTestData( 137 | [{"func_name": "nop", 138 | "path": "", 139 | "line": 1, 140 | "signature": { 141 | "arg_types": ["int"], 142 | "return_type": "object"}, 143 | }]) 144 | a = """\ 145 | def nop(foo, **kwds): 146 | return 42 147 | """ 148 | b = """\ 149 | from typing import Any 150 | def nop(foo: int, **kwds: Any) -> object: 151 | return 42 152 | """ 153 | self.check(a, b) 154 | 155 | def test_dont_add_kwds(self): 156 | self.setTestData( 157 | [{"func_name": "nop", 158 | "path": "", 159 | "line": 1, 160 | "signature": { 161 | "arg_types": ["int", "**AnyStr"], 162 | "return_type": "object"}, 163 | }]) 164 | a = """\ 165 | def nop(foo, **kwds): 166 | return 42 167 | """ 168 | b = """\ 169 | from typing import AnyStr 170 | def nop(foo: int, **kwds: AnyStr) -> object: 171 | return 42 172 | """ 173 | self.check(a, b) 174 | 175 | def test_add_varargs(self): 176 | self.setTestData( 177 | [{"func_name": "nop", 178 | "path": "", 179 | "line": 1, 180 | "signature": { 181 | "arg_types": ["int"], 182 | "return_type": "object"}, 183 | }]) 184 | a = """\ 185 | def nop(foo, *args): 186 | return 42 187 | """ 188 | b = """\ 189 | from typing import Any 190 | def nop(foo: int, *args: Any) -> object: 191 | return 42 192 | """ 193 | self.check(a, b) 194 | 195 | def test_dont_add_varargs(self): 196 | self.setTestData( 197 | [{"func_name": "nop", 198 | "path": "", 199 | "line": 1, 200 | "signature": { 201 | "arg_types": ["int", "*int"], 202 | "return_type": "object"}, 203 | }]) 204 | a = """\ 205 | def nop(foo, *args): 206 | return 42 207 | """ 208 | b = """\ 209 | def nop(foo: int, *args: int) -> object: 210 | return 42 211 | """ 212 | self.check(a, b) 213 | 214 | def test_return_expr_not_none(self): 215 | self.setTestData( 216 | [{"func_name": "nop", 217 | "path": "", 218 | "line": 1, 219 | "signature": { 220 | "arg_types": [], 221 | "return_type": "None"}, 222 | }]) 223 | a = """\ 224 | def nop(): 225 | return 0 226 | """ 227 | b = """\ 228 | from typing import Any 229 | from typing import Optional 230 | def nop() -> Optional[Any]: 231 | return 0 232 | """ 233 | self.check(a, b) 234 | 235 | def test_return_expr_none(self): 236 | self.setTestData( 237 | [{"func_name": "nop", 238 | "path": "", 239 | "line": 1, 240 | "signature": { 241 | "arg_types": [], 242 | "return_type": "None"}, 243 | }]) 244 | a = """\ 245 | def nop(): 246 | return 247 | """ 248 | b = """\ 249 | def nop() -> None: 250 | return 251 | """ 252 | self.check(a, b) 253 | 254 | def test_generator_optional(self): 255 | self.setTestData( 256 | [{"func_name": "gen", 257 | "path": "", 258 | "line": 1, 259 | "signature": { 260 | "arg_types": [], 261 | "return_type": "Optional[int]"}, 262 | }]) 263 | a = """\ 264 | def gen(): 265 | yield 42 266 | """ 267 | b = """\ 268 | from typing import Iterator 269 | def gen() -> Iterator[int]: 270 | yield 42 271 | """ 272 | self.check(a, b) 273 | 274 | def test_generator_plain(self): 275 | self.setTestData( 276 | [{"func_name": "gen", 277 | "path": "", 278 | "line": 1, 279 | "signature": { 280 | "arg_types": [], 281 | "return_type": "int"}, 282 | }]) 283 | a = """\ 284 | def gen(): 285 | yield 42 286 | """ 287 | b = """\ 288 | from typing import Iterator 289 | def gen() -> Iterator[int]: 290 | yield 42 291 | """ 292 | self.check(a, b) 293 | 294 | def test_not_generator(self): 295 | self.setTestData( 296 | [{"func_name": "nop", 297 | "path": "", 298 | "line": 1, 299 | "signature": { 300 | "arg_types": [], 301 | "return_type": "int"}, 302 | }]) 303 | a = """\ 304 | def nop(): 305 | def gen(): 306 | yield 42 307 | """ 308 | b = """\ 309 | def nop() -> int: 310 | def gen(): 311 | yield 42 312 | """ 313 | self.check(a, b) 314 | 315 | def test_add_self(self): 316 | self.setTestData( 317 | [{"func_name": "nop", 318 | "path": "", 319 | "line": 1, 320 | "signature": { 321 | "arg_types": [], 322 | "return_type": "int"}, 323 | }]) 324 | a = """\ 325 | def nop(self): 326 | pass 327 | """ 328 | b = """\ 329 | from typing import Any 330 | def nop(self: Any) -> int: 331 | pass 332 | """ 333 | self.check(a, b) 334 | 335 | def test_dont_add_self(self): 336 | self.setTestData( 337 | [{"func_name": "C.nop", 338 | "path": "", 339 | "line": 1, 340 | "signature": { 341 | "arg_types": [], 342 | "return_type": "int"}, 343 | }]) 344 | a = """\ 345 | class C: 346 | def nop(self): 347 | pass 348 | """ 349 | b = """\ 350 | class C: 351 | def nop(self) -> int: 352 | pass 353 | """ 354 | self.check(a, b) 355 | 356 | def test_too_many_types(self): 357 | self.setTestData( 358 | [{"func_name": "nop", 359 | "path": "", 360 | "line": 1, 361 | "signature": { 362 | "arg_types": ["int"], 363 | "return_type": "int"}, 364 | }]) 365 | a = """\ 366 | def nop(): 367 | pass 368 | """ 369 | self.warns(a, a, "source has 0 args, annotation has 1 -- skipping", unchanged=True) 370 | 371 | def test_too_few_types(self): 372 | self.setTestData( 373 | [{"func_name": "nop", 374 | "path": "", 375 | "line": 1, 376 | "signature": { 377 | "arg_types": [], 378 | "return_type": "int"}, 379 | }]) 380 | a = """\ 381 | def nop(a): 382 | pass 383 | """ 384 | self.warns(a, a, "source has 1 args, annotation has 0 -- skipping", unchanged=True) 385 | 386 | def test_line_number_drift(self): 387 | self.setTestData( 388 | [{"func_name": "nop", 389 | "path": "", 390 | "line": 10, 391 | "signature": { 392 | "arg_types": [], 393 | "return_type": "int"}, 394 | }]) 395 | a = """\ 396 | def nop(a): 397 | pass 398 | """ 399 | self.warns(a, a, "signature from line 10 too far away -- skipping", unchanged=True) 400 | 401 | def test_classmethod(self): 402 | # Class method names currently are returned without class name 403 | self.setTestData( 404 | [{"func_name": "nop", 405 | "path": "", 406 | "line": 3, 407 | "signature": { 408 | "arg_types": ["int"], 409 | "return_type": "int"} 410 | }]) 411 | a = """\ 412 | class C: 413 | @classmethod 414 | def nop(cls, a): 415 | return a 416 | """ 417 | b = """\ 418 | class C: 419 | @classmethod 420 | def nop(cls, a: int) -> int: 421 | return a 422 | """ 423 | self.check(a, b) 424 | 425 | def test_staticmethod(self): 426 | # Static method names currently are returned without class name 427 | self.setTestData( 428 | [{"func_name": "nop", 429 | "path": "", 430 | "line": 3, 431 | "signature": { 432 | "arg_types": ["int"], 433 | "return_type": "int"} 434 | }]) 435 | a = """\ 436 | class C: 437 | @staticmethod 438 | def nop(a): 439 | return a 440 | """ 441 | b = """\ 442 | class C: 443 | @staticmethod 444 | def nop(a: int) -> int: 445 | return a 446 | """ 447 | self.check(a, b) 448 | 449 | def test_long_form(self): 450 | self.maxDiff = None 451 | self.setTestData( 452 | [{"func_name": "nop", 453 | "path": "", 454 | "line": 1, 455 | "signature": { 456 | "arg_types": ["int", "int", "int", 457 | "str", "str", "str", 458 | "Optional[bool]", "Union[int, str]", "*Any"], 459 | "return_type": "int"}, 460 | }]) 461 | a = """\ 462 | def nop(a, b, c, # some comment 463 | d, e, f, # multi-line 464 | # comment 465 | g=None, h=0, *args): 466 | return 0 467 | """ 468 | b = """\ 469 | from typing import Any 470 | from typing import Optional 471 | from typing import Union 472 | def nop(a: int, b: int, c: int, # some comment 473 | d: str, e: str, f: str, # multi-line 474 | # comment 475 | g: Optional[bool] = None, h: Union[int, str] = 0, *args: Any) -> int: 476 | return 0 477 | """ 478 | self.check(a, b) 479 | 480 | def test_long_form_method(self): 481 | self.maxDiff = None 482 | self.setTestData( 483 | [{"func_name": "C.nop", 484 | "path": "", 485 | "line": 2, 486 | "signature": { 487 | "arg_types": ["int", "int", "int", 488 | "str", "str", "str", 489 | "Optional[bool]", "Union[int, str]", "*Any"], 490 | "return_type": "int"}, 491 | }]) 492 | a = """\ 493 | class C: 494 | def nop(self, a, b, c, # some comment 495 | d, e, f, # multi-line 496 | # comment 497 | g=None, h=0, *args): 498 | return 0 499 | """ 500 | b = """\ 501 | from typing import Any 502 | from typing import Optional 503 | from typing import Union 504 | class C: 505 | def nop(self, a: int, b: int, c: int, # some comment 506 | d: str, e: str, f: str, # multi-line 507 | # comment 508 | g: Optional[bool] = None, h: Union[int, str] = 0, *args: Any) -> int: 509 | return 0 510 | """ 511 | self.check(a, b) 512 | 513 | def test_long_form_classmethod(self): 514 | self.maxDiff = None 515 | self.setTestData( 516 | [{"func_name": "nop", 517 | "path": "", 518 | "line": 3, 519 | "signature": { 520 | "arg_types": ["int", "int", "int", 521 | "str", "str", "str", 522 | "Optional[bool]", "Union[int, str]", "*Any"], 523 | "return_type": "int"}, 524 | }]) 525 | a = """\ 526 | class C: 527 | @classmethod 528 | def nop(cls, a, b, c, # some comment 529 | d, e, f, 530 | g=None, h=0, *args): 531 | return 0 532 | """ 533 | b = """\ 534 | from typing import Any 535 | from typing import Optional 536 | from typing import Union 537 | class C: 538 | @classmethod 539 | def nop(cls, a: int, b: int, c: int, # some comment 540 | d: str, e: str, f: str, 541 | g: Optional[bool] = None, h: Union[int, str] = 0, *args: Any) -> int: 542 | return 0 543 | """ 544 | self.check(a, b) 545 | # Do the same test for staticmethod 546 | a = a.replace('classmethod', 'staticmethod') 547 | b = b.replace('classmethod', 'staticmethod') 548 | self.check(a, b) 549 | 550 | def test_long_form_trailing_comma(self): 551 | self.maxDiff = None 552 | self.setTestData( 553 | [{"func_name": "nop", 554 | "path": "", 555 | "line": 3, 556 | "signature": { 557 | "arg_types": ["int", "int", "int", 558 | "str", "str", "str", 559 | "Optional[bool]", "Union[int, str]"], 560 | "return_type": "int"}, 561 | }]) 562 | a = """\ 563 | def nop(a, b, c, # some comment 564 | d, e, f, 565 | g=None, h=0): 566 | return 0 567 | """ 568 | b = """\ 569 | from typing import Optional 570 | from typing import Union 571 | def nop(a: int, b: int, c: int, # some comment 572 | d: str, e: str, f: str, 573 | g: Optional[bool] = None, h: Union[int, str] = 0) -> int: 574 | return 0 575 | """ 576 | self.check(a, b) 577 | 578 | def test_one_liner(self): 579 | self.setTestData( 580 | [{"func_name": "nop", 581 | "path": "", 582 | "line": 1, 583 | "signature": { 584 | "arg_types": ["int"], 585 | "return_type": "int"}, 586 | }]) 587 | a = """\ 588 | def nop(a): return a 589 | """ 590 | b = """\ 591 | def nop(a: int) -> int: return a 592 | """ 593 | self.check(a, b) 594 | 595 | def test_variadic(self): 596 | self.setTestData( 597 | [{"func_name": "nop", 598 | "path": "", 599 | "line": 1, 600 | "signature": { 601 | "arg_types": ["Tuple[int, ...]"], 602 | "return_type": "int"}, 603 | }]) 604 | a = """\ 605 | def nop(a): return 0 606 | """ 607 | b = """\ 608 | from typing import Tuple 609 | def nop(a: Tuple[int, ...]) -> int: return 0 610 | """ 611 | self.check(a, b) 612 | 613 | @unittest.skipIf(sys.version_info < (3, 5), 'async not supported on old python') 614 | def test_nested_class_async_func(self): 615 | self.setTestData( 616 | [{"func_name": "A.B.foo", 617 | "path": "", 618 | "line": 3, 619 | "signature": { 620 | "arg_types": ['str'], 621 | "return_type": "int"}, 622 | }]) 623 | a = """\ 624 | class A: 625 | class B: 626 | async def foo(x): 627 | return 42 628 | """ 629 | b = """\ 630 | class A: 631 | class B: 632 | async def foo(x: str) -> int: 633 | return 42 634 | """ 635 | self.check(a, b) 636 | 637 | @patch('pyannotate_tools.fixes.fix_annotate_json.FixAnnotateJson.set_filename') 638 | def test_set_filename(self, mocked_set_filename): 639 | self.filename = "/path/to/fileA.py" 640 | self.unchanged("") 641 | mocked_set_filename.assert_called_with("/path/to/fileA.py") 642 | 643 | self.filename = "/path/to/fileB.py" 644 | self.unchanged("") 645 | mocked_set_filename.assert_called_with("/path/to/fileB.py") 646 | -------------------------------------------------------------------------------- /pyannotate_tools/fixes/tests/test_annotate_py2.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | # Our flake extension misfires on type comments in strings below. 3 | 4 | from lib2to3.tests.test_fixers import FixerTestCase 5 | 6 | # deadcode: fix_annotate is used as part of the fixer_pkg for this test 7 | from pyannotate_tools.fixes import fix_annotate 8 | 9 | 10 | class TestFixAnnotate(FixerTestCase): 11 | 12 | def setUp(self): 13 | super(TestFixAnnotate, self).setUp( 14 | fix_list=["annotate"], 15 | fixer_pkg="pyannotate_tools", 16 | options={'annotation_style' : 'py2'}, 17 | ) 18 | 19 | def test_no_arg(self): 20 | a = """\ 21 | def nop(): 22 | return 42 23 | """ 24 | b = """\ 25 | from typing import Any 26 | def nop(): 27 | # type: () -> Any 28 | return 42 29 | """ 30 | self.check(a, b) 31 | 32 | def test_one_arg(self): 33 | a = """\ 34 | def incr(arg): 35 | return arg+1 36 | """ 37 | b = """\ 38 | from typing import Any 39 | def incr(arg): 40 | # type: (Any) -> Any 41 | return arg+1 42 | """ 43 | self.check(a, b) 44 | 45 | def test_two_args(self): 46 | a = """\ 47 | def add(arg1, arg2): 48 | return arg1+arg2 49 | """ 50 | b = """\ 51 | from typing import Any 52 | def add(arg1, arg2): 53 | # type: (Any, Any) -> Any 54 | return arg1+arg2 55 | """ 56 | self.check(a, b) 57 | 58 | def test_defaults(self): 59 | a = """\ 60 | def foo(iarg=0, farg=0.0, sarg='', uarg=u'', barg=False): 61 | return 42 62 | """ 63 | b = """\ 64 | from typing import Any 65 | def foo(iarg=0, farg=0.0, sarg='', uarg=u'', barg=False): 66 | # type: (int, float, str, unicode, bool) -> Any 67 | return 42 68 | """ 69 | self.check(a, b) 70 | 71 | def test_staticmethod(self): 72 | a = """\ 73 | class C: 74 | @staticmethod 75 | def incr(self): 76 | return 42 77 | """ 78 | b = """\ 79 | from typing import Any 80 | class C: 81 | @staticmethod 82 | def incr(self): 83 | # type: (Any) -> Any 84 | return 42 85 | """ 86 | self.check(a, b) 87 | 88 | def test_classmethod(self): 89 | a = """\ 90 | class C: 91 | @classmethod 92 | def incr(cls, arg): 93 | return 42 94 | """ 95 | b = """\ 96 | from typing import Any 97 | class C: 98 | @classmethod 99 | def incr(cls, arg): 100 | # type: (Any) -> Any 101 | return 42 102 | """ 103 | self.check(a, b) 104 | 105 | def test_instancemethod(self): 106 | a = """\ 107 | class C: 108 | def incr(self, arg): 109 | return 42 110 | """ 111 | b = """\ 112 | from typing import Any 113 | class C: 114 | def incr(self, arg): 115 | # type: (Any) -> Any 116 | return 42 117 | """ 118 | self.check(a, b) 119 | 120 | def test_fake_self(self): 121 | a = """\ 122 | def incr(self, arg): 123 | return 42 124 | """ 125 | b = """\ 126 | from typing import Any 127 | def incr(self, arg): 128 | # type: (Any, Any) -> Any 129 | return 42 130 | """ 131 | self.check(a, b) 132 | 133 | def test_nested_fake_self(self): 134 | a = """\ 135 | class C: 136 | def outer(self): 137 | def inner(self, arg): 138 | return 42 139 | """ 140 | b = """\ 141 | from typing import Any 142 | class C: 143 | def outer(self): 144 | # type: () -> None 145 | def inner(self, arg): 146 | # type: (Any, Any) -> Any 147 | return 42 148 | """ 149 | self.check(a, b) 150 | 151 | def test_multiple_decorators(self): 152 | a = """\ 153 | class C: 154 | @contextmanager 155 | @classmethod 156 | @wrapped('func') 157 | def incr(cls, arg): 158 | return 42 159 | """ 160 | b = """\ 161 | from typing import Any 162 | class C: 163 | @contextmanager 164 | @classmethod 165 | @wrapped('func') 166 | def incr(cls, arg): 167 | # type: (Any) -> Any 168 | return 42 169 | """ 170 | self.check(a, b) 171 | 172 | def test_stars(self): 173 | a = """\ 174 | def stuff(*a, **kw): 175 | return 4, 2 176 | """ 177 | b = """\ 178 | from typing import Any 179 | def stuff(*a, **kw): 180 | # type: (*Any, **Any) -> Any 181 | return 4, 2 182 | """ 183 | self.check(a, b) 184 | 185 | def test_idempotency(self): 186 | a = """\ 187 | def incr(arg): 188 | # type: (Any) -> Any 189 | return arg+1 190 | """ 191 | self.unchanged(a) 192 | 193 | def test_no_return_expr(self): 194 | a = """\ 195 | def proc1(arg): 196 | return 197 | def proc2(arg): 198 | pass 199 | """ 200 | b = """\ 201 | from typing import Any 202 | def proc1(arg): 203 | # type: (Any) -> None 204 | return 205 | def proc2(arg): 206 | # type: (Any) -> None 207 | pass 208 | """ 209 | self.check(a, b) 210 | 211 | def test_nested_return_expr(self): 212 | # The 'return expr' in inner() shouldn't affect the return type of outer(). 213 | a = """\ 214 | def outer(arg): 215 | def inner(): 216 | return 42 217 | return 218 | """ 219 | b = """\ 220 | from typing import Any 221 | def outer(arg): 222 | # type: (Any) -> None 223 | def inner(): 224 | # type: () -> Any 225 | return 42 226 | return 227 | """ 228 | self.check(a, b) 229 | 230 | def test_nested_class_return_expr(self): 231 | # The 'return expr' in class Inner shouldn't affect the return type of outer(). 232 | a = """\ 233 | def outer(arg): 234 | class Inner: 235 | return 42 236 | return 237 | """ 238 | b = """\ 239 | from typing import Any 240 | def outer(arg): 241 | # type: (Any) -> None 242 | class Inner: 243 | return 42 244 | return 245 | """ 246 | self.check(a, b) 247 | 248 | def test_add_import(self): 249 | a = """\ 250 | import typing 251 | from typing import Callable 252 | 253 | def incr(arg): 254 | return 42 255 | """ 256 | b = """\ 257 | import typing 258 | from typing import Callable 259 | from typing import Any 260 | 261 | def incr(arg): 262 | # type: (Any) -> Any 263 | return 42 264 | """ 265 | self.check(a, b) 266 | 267 | def test_dont_add_import(self): 268 | a = """\ 269 | def nop(arg=0): 270 | return 271 | """ 272 | b = """\ 273 | def nop(arg=0): 274 | # type: (int) -> None 275 | return 276 | """ 277 | self.check(a, b) 278 | 279 | def test_long_form(self): 280 | self.maxDiff = None 281 | a = """\ 282 | def nop(arg0, arg1, arg2, arg3, arg4, 283 | arg5, arg6, arg7, arg8=0, arg9='', 284 | *args, **kwds): 285 | return 286 | """ 287 | b = """\ 288 | from typing import Any 289 | def nop(arg0, # type: Any 290 | arg1, # type: Any 291 | arg2, # type: Any 292 | arg3, # type: Any 293 | arg4, # type: Any 294 | arg5, # type: Any 295 | arg6, # type: Any 296 | arg7, # type: Any 297 | arg8=0, # type: int 298 | arg9='', # type: str 299 | *args, # type: Any 300 | **kwds # type: Any 301 | ): 302 | # type: (...) -> None 303 | return 304 | """ 305 | self.check(a, b) 306 | 307 | def test_long_form_trailing_comma(self): 308 | self.maxDiff = None 309 | a = """\ 310 | def nop(arg0, arg1, arg2, arg3, arg4, arg5, arg6, 311 | arg7=None, arg8=0, arg9='', arg10=False,): 312 | return 313 | """ 314 | b = """\ 315 | from typing import Any 316 | def nop(arg0, # type: Any 317 | arg1, # type: Any 318 | arg2, # type: Any 319 | arg3, # type: Any 320 | arg4, # type: Any 321 | arg5, # type: Any 322 | arg6, # type: Any 323 | arg7=None, # type: Any 324 | arg8=0, # type: int 325 | arg9='', # type: str 326 | arg10=False, # type: bool 327 | ): 328 | # type: (...) -> None 329 | return 330 | """ 331 | self.check(a, b) 332 | 333 | def test_one_liner(self): 334 | a = """\ 335 | class C: 336 | def nop(self, a): a = a; return a 337 | # Something 338 | # More 339 | pass 340 | """ 341 | b = """\ 342 | from typing import Any 343 | class C: 344 | def nop(self, a): 345 | # type: (Any) -> Any 346 | a = a; return a 347 | # Something 348 | # More 349 | pass 350 | """ 351 | self.check(a, b) 352 | 353 | def test_idempotency_long_1arg(self): 354 | a = """\ 355 | def nop(a # type: int 356 | ): 357 | pass 358 | """ 359 | self.unchanged(a) 360 | 361 | def test_idempotency_long_1arg_comma(self): 362 | a = """\ 363 | def nop(a, # type: int 364 | ): 365 | pass 366 | """ 367 | self.unchanged(a) 368 | 369 | def test_idempotency_long_2args_first(self): 370 | a = """\ 371 | def nop(a, # type: int 372 | b): 373 | pass 374 | """ 375 | self.unchanged(a) 376 | 377 | def test_idempotency_long_2args_last(self): 378 | a = """\ 379 | def nop(a, 380 | b # type: int 381 | ): 382 | pass 383 | """ 384 | self.unchanged(a) 385 | 386 | def test_idempotency_long_varargs(self): 387 | a = """\ 388 | def nop(*a # type: int 389 | ): 390 | pass 391 | """ 392 | self.unchanged(a) 393 | 394 | def test_idempotency_long_kwargs(self): 395 | a = """\ 396 | def nop(**a # type: int 397 | ): 398 | pass 399 | """ 400 | self.unchanged(a) 401 | -------------------------------------------------------------------------------- /pyannotate_tools/fixes/tests/test_annotate_py3.py: -------------------------------------------------------------------------------- 1 | # flake8: noqa 2 | # Our flake extension misfires on type comments in strings below. 3 | 4 | from lib2to3.tests.test_fixers import FixerTestCase 5 | import unittest 6 | 7 | # deadcode: fix_annotate is used as part of the fixer_pkg for this test 8 | from pyannotate_tools.fixes import fix_annotate 9 | 10 | 11 | class TestFixAnnotate3(FixerTestCase): 12 | 13 | def setUp(self): 14 | super(TestFixAnnotate3, self).setUp( 15 | fix_list=["annotate"], 16 | fixer_pkg="pyannotate_tools", 17 | options={'annotation_style' : 'py3'} 18 | ) 19 | 20 | def test_no_arg_1(self) : 21 | a = """\ 22 | def nop(): 23 | return 42 24 | """ 25 | b = """\ 26 | from typing import Any 27 | def nop() -> Any: 28 | return 42 29 | """ 30 | self.check(a, b) 31 | 32 | def test_no_arg_2(self) : 33 | a = """\ 34 | def nop(): return 42 35 | """ 36 | b = """\ 37 | from typing import Any 38 | def nop() -> Any: return 42 39 | """ 40 | self.check(a, b) 41 | 42 | def test_no_arg_3(self) : 43 | a = """\ 44 | def nop( 45 | ): 46 | return 42 47 | """ 48 | b = """\ 49 | from typing import Any 50 | def nop( 51 | ) -> Any: 52 | return 42 53 | """ 54 | self.check(a, b) 55 | 56 | def test_no_arg_4(self) : 57 | a = """\ 58 | def nop( 59 | ) \ 60 | : 61 | return 42 62 | """ 63 | b = """\ 64 | from typing import Any 65 | def nop( 66 | ) -> Any \ 67 | : 68 | return 42 69 | """ 70 | self.check(a, b) 71 | 72 | def test_no_arg_5(self) : 73 | a = """\ 74 | def nop( # blah 75 | ): # blah 76 | return 42 # blah 77 | """ 78 | b = """\ 79 | from typing import Any 80 | def nop( # blah 81 | ) -> Any: # blah 82 | return 42 # blah 83 | """ 84 | self.check(a, b) 85 | 86 | def test_no_arg_6(self) : 87 | a = """\ 88 | def nop( # blah 89 | ) \ 90 | : # blah 91 | return 42 # blah 92 | """ 93 | b = """\ 94 | from typing import Any 95 | def nop( # blah 96 | ) -> Any \ 97 | : # blah 98 | return 42 # blah 99 | """ 100 | self.check(a, b) 101 | 102 | def test_one_arg_1(self): 103 | a = """\ 104 | def incr(arg): 105 | return arg+1 106 | """ 107 | b = """\ 108 | from typing import Any 109 | def incr(arg: Any) -> Any: 110 | return arg+1 111 | """ 112 | self.check(a, b) 113 | 114 | 115 | def test_one_arg_2(self): 116 | a = """\ 117 | def incr(arg=0): 118 | return arg+1 119 | """ 120 | b = """\ 121 | from typing import Any 122 | def incr(arg: int = 0) -> Any: 123 | return arg+1 124 | """ 125 | self.check(a, b) 126 | 127 | def test_one_arg_3(self): 128 | a = """\ 129 | def incr( arg=0 ): 130 | return arg+1 131 | """ 132 | b = """\ 133 | from typing import Any 134 | def incr( arg: int = 0 ) -> Any: 135 | return arg+1 136 | """ 137 | self.check(a, b) 138 | 139 | def test_one_arg_4(self): 140 | a = """\ 141 | def incr( arg = 0 ): 142 | return arg+1 143 | """ 144 | b = """\ 145 | from typing import Any 146 | def incr( arg: int = 0 ) -> Any: 147 | return arg+1 148 | """ 149 | self.check(a, b) 150 | 151 | def test_two_args_1(self): 152 | a = """\ 153 | def add(arg1, arg2): 154 | return arg1+arg2 155 | """ 156 | b = """\ 157 | from typing import Any 158 | def add(arg1: Any, arg2: Any) -> Any: 159 | return arg1+arg2 160 | """ 161 | self.check(a, b) 162 | 163 | def test_two_args_2(self): 164 | a = """\ 165 | def add(arg1=0, arg2=0.1): 166 | return arg1+arg2 167 | """ 168 | b = """\ 169 | from typing import Any 170 | def add(arg1: int = 0, arg2: float = 0.1) -> Any: 171 | return arg1+arg2 172 | """ 173 | self.check(a, b) 174 | 175 | def test_two_args_3(self): 176 | a = """\ 177 | def add(arg1, arg2=0.1): 178 | return arg1+arg2 179 | """ 180 | b = """\ 181 | from typing import Any 182 | def add(arg1: Any, arg2: float = 0.1) -> Any: 183 | return arg1+arg2 184 | """ 185 | 186 | def test_two_args_4(self): 187 | a = """\ 188 | def add(arg1, arg2 = 0.1): 189 | return arg1+arg2 190 | """ 191 | b = """\ 192 | from typing import Any 193 | def add(arg1: Any, arg2: float = 0.1) -> Any: 194 | return arg1+arg2 195 | """ 196 | self.check(a, b) 197 | self.check(a, b) 198 | 199 | def test_defaults_1(self): 200 | a = """\ 201 | def foo(iarg=0, farg=0.0, sarg='', uarg=u'', barg=False): 202 | return 42 203 | """ 204 | b = """\ 205 | from typing import Any 206 | def foo(iarg: int = 0, farg: float = 0.0, sarg: str = '', uarg: unicode = u'', barg: bool = False) -> Any: 207 | return 42 208 | """ 209 | self.check(a, b) 210 | 211 | def test_defaults_2(self): 212 | a = """\ 213 | def foo(iarg=0, farg=0.0, sarg='', uarg=u'', barg=False, targ=(1,2,3)): 214 | return 42 215 | """ 216 | b = """\ 217 | from typing import Any 218 | def foo(iarg: int = 0, farg: float = 0.0, sarg: str = '', uarg: unicode = u'', barg: bool = False, targ: Any = (1,2,3)) -> Any: 219 | return 42 220 | """ 221 | self.check(a, b) 222 | 223 | def test_defaults_3(self): 224 | a = """\ 225 | def foo(iarg=0, farg, sarg='', uarg, barg=False, targ=(1,2,3)): 226 | return 42 227 | """ 228 | b = """\ 229 | from typing import Any 230 | def foo(iarg: int = 0, farg: Any, sarg: str = '', uarg: Any, barg: bool = False, targ: Any = (1,2,3)) -> Any: 231 | return 42 232 | """ 233 | self.check(a, b) 234 | 235 | def test_staticmethod(self): 236 | a = """\ 237 | class C: 238 | @staticmethod 239 | def incr(self): 240 | return 42 241 | """ 242 | b = """\ 243 | from typing import Any 244 | class C: 245 | @staticmethod 246 | def incr(self: Any) -> Any: 247 | return 42 248 | """ 249 | self.check(a, b) 250 | 251 | def test_classmethod(self): 252 | a = """\ 253 | class C: 254 | @classmethod 255 | def incr(cls, arg): 256 | return 42 257 | """ 258 | b = """\ 259 | from typing import Any 260 | class C: 261 | @classmethod 262 | def incr(cls, arg: Any) -> Any: 263 | return 42 264 | """ 265 | self.check(a, b) 266 | 267 | def test_instancemethod(self): 268 | a = """\ 269 | class C: 270 | def incr(self, arg): 271 | return 42 272 | """ 273 | b = """\ 274 | from typing import Any 275 | class C: 276 | def incr(self, arg: Any) -> Any: 277 | return 42 278 | """ 279 | self.check(a, b) 280 | 281 | def test_fake_self(self): 282 | a = """\ 283 | def incr(self, arg): 284 | return 42 285 | """ 286 | b = """\ 287 | from typing import Any 288 | def incr(self: Any, arg: Any) -> Any: 289 | return 42 290 | """ 291 | self.check(a, b) 292 | 293 | def test_nested_fake_self(self): 294 | a = """\ 295 | class C: 296 | def outer(self): 297 | def inner(self, arg): 298 | return 42 299 | """ 300 | b = """\ 301 | from typing import Any 302 | class C: 303 | def outer(self) -> None: 304 | def inner(self: Any, arg: Any) -> Any: 305 | return 42 306 | """ 307 | self.check(a, b) 308 | 309 | def test_multiple_decorators(self): 310 | a = """\ 311 | class C: 312 | @contextmanager 313 | @classmethod 314 | @wrapped('func') 315 | def incr(cls, arg): 316 | return 42 317 | """ 318 | b = """\ 319 | from typing import Any 320 | class C: 321 | @contextmanager 322 | @classmethod 323 | @wrapped('func') 324 | def incr(cls, arg: Any) -> Any: 325 | return 42 326 | """ 327 | self.check(a, b) 328 | 329 | def test_stars_1(self): 330 | a = """\ 331 | def stuff(*a): 332 | return 4, 2 333 | """ 334 | b = """\ 335 | from typing import Any 336 | def stuff(*a: Any) -> Any: 337 | return 4, 2 338 | """ 339 | self.check(a, b) 340 | 341 | def test_stars_2(self): 342 | a = """\ 343 | def stuff(a, *b): 344 | return 4, 2 345 | """ 346 | b = """\ 347 | from typing import Any 348 | def stuff(a: Any, *b: Any) -> Any: 349 | return 4, 2 350 | """ 351 | self.check(a, b) 352 | 353 | 354 | def test_keywords_1(self): 355 | a = """\ 356 | def stuff(**kw): 357 | return 4, 2 358 | """ 359 | b = """\ 360 | from typing import Any 361 | def stuff(**kw: Any) -> Any: 362 | return 4, 2 363 | """ 364 | self.check(a, b) 365 | 366 | def test_keywords_2(self): 367 | a = """\ 368 | def stuff(a, **kw): 369 | return 4, 2 370 | """ 371 | b = """\ 372 | from typing import Any 373 | def stuff(a: Any, **kw: Any) -> Any: 374 | return 4, 2 375 | """ 376 | self.check(a, b) 377 | 378 | def test_keywords_3(self): 379 | a = """\ 380 | def stuff(a, *b, **kw): 381 | return 4, 2 382 | """ 383 | b = """\ 384 | from typing import Any 385 | def stuff(a: Any, *b: Any, **kw: Any) -> Any: 386 | return 4, 2 387 | """ 388 | self.check(a, b) 389 | 390 | def test_keywords_4(self): 391 | a = """\ 392 | def stuff(*b, **kw): 393 | return 4, 2 394 | """ 395 | b = """\ 396 | from typing import Any 397 | def stuff(*b: Any, **kw: Any) -> Any: 398 | return 4, 2 399 | """ 400 | self.check(a, b) 401 | 402 | def test_no_return_expr(self): 403 | a = """\ 404 | def proc1(arg): 405 | return 406 | def proc2(arg): 407 | pass 408 | """ 409 | b = """\ 410 | from typing import Any 411 | def proc1(arg: Any) -> None: 412 | return 413 | def proc2(arg: Any) -> None: 414 | pass 415 | """ 416 | self.check(a, b) 417 | 418 | def test_nested_return_expr(self): 419 | # The 'return expr' in inner() shouldn't affect the return type of outer(). 420 | a = """\ 421 | def outer(arg): 422 | def inner(): 423 | return 42 424 | return 425 | """ 426 | b = """\ 427 | from typing import Any 428 | def outer(arg: Any) -> None: 429 | def inner() -> Any: 430 | return 42 431 | return 432 | """ 433 | self.check(a, b) 434 | 435 | def test_nested_class_return_expr(self): 436 | # The 'return expr' in class Inner shouldn't affect the return type of outer(). 437 | a = """\ 438 | def outer(arg): 439 | class Inner: 440 | return 42 441 | return 442 | """ 443 | b = """\ 444 | from typing import Any 445 | def outer(arg: Any) -> None: 446 | class Inner: 447 | return 42 448 | return 449 | """ 450 | self.check(a, b) 451 | 452 | def test_add_import(self): 453 | a = """\ 454 | import typing 455 | from typing import Callable 456 | 457 | def incr(arg): 458 | return 42 459 | """ 460 | b = """\ 461 | import typing 462 | from typing import Callable 463 | from typing import Any 464 | 465 | def incr(arg: Any) -> Any: 466 | return 42 467 | """ 468 | self.check(a, b) 469 | 470 | def test_dont_add_import(self): 471 | a = """\ 472 | def nop(arg=0): 473 | return 474 | """ 475 | b = """\ 476 | def nop(arg: int = 0) -> None: 477 | return 478 | """ 479 | self.check(a, b) 480 | 481 | def test_long_form(self): 482 | self.maxDiff = None 483 | a = """\ 484 | def nop(arg0, arg1, arg2, arg3, arg4, 485 | arg5, arg6, arg7, arg8=0, arg9='', 486 | *args, **kwds): 487 | return 488 | """ 489 | b = """\ 490 | from typing import Any 491 | def nop(arg0: Any, arg1: Any, arg2: Any, arg3: Any, arg4: Any, 492 | arg5: Any, arg6: Any, arg7: Any, arg8: int = 0, arg9: str = '', 493 | *args: Any, **kwds: Any) -> None: 494 | return 495 | """ 496 | self.check(a, b) 497 | 498 | def test_long_form_trailing_comma(self): 499 | self.maxDiff = None 500 | a = """\ 501 | def nop(arg0, arg1, arg2, arg3, arg4, arg5, arg6, 502 | arg7=None, arg8=0, arg9='', arg10=False,): 503 | return 504 | """ 505 | b = """\ 506 | from typing import Any 507 | def nop(arg0: Any, arg1: Any, arg2: Any, arg3: Any, arg4: Any, arg5: Any, arg6: Any, 508 | arg7: Any = None, arg8: int = 0, arg9: str = '', arg10: bool = False,) -> None: 509 | return 510 | """ 511 | self.check(a, b) 512 | 513 | def test_one_liner(self): 514 | a = """\ 515 | class C: 516 | def nop(self, a): a = a; return a 517 | # Something 518 | # More 519 | pass 520 | """ 521 | b = """\ 522 | from typing import Any 523 | class C: 524 | def nop(self, a: Any) -> Any: a = a; return a 525 | # Something 526 | # More 527 | pass 528 | """ 529 | self.check(a, b) 530 | 531 | def test_idempotency_long_1arg(self): 532 | a = """\ 533 | def nop(a: int 534 | ): 535 | pass 536 | """ 537 | self.unchanged(a) 538 | 539 | def test_idempotency_long_1arg_comma(self): 540 | a = """\ 541 | def nop(a: int, 542 | ): 543 | pass 544 | """ 545 | self.unchanged(a) 546 | 547 | def test_idempotency_long_2args_first(self): 548 | a = """\ 549 | def nop(a: int, 550 | b): 551 | pass 552 | """ 553 | self.unchanged(a) 554 | 555 | def test_idempotency_long_2args_last(self): 556 | a = """\ 557 | def nop(a, 558 | b: int 559 | ): 560 | pass 561 | """ 562 | self.unchanged(a) 563 | 564 | def test_idempotency_long_varargs(self): 565 | a = """\ 566 | def nop(*a: int 567 | ): 568 | pass 569 | """ 570 | self.unchanged(a) 571 | 572 | def test_idempotency_long_kwargs(self): 573 | a = """\ 574 | def nop(**a: int 575 | ): 576 | pass 577 | """ 578 | self.unchanged(a) 579 | 580 | 581 | def test_idempotency_arg0_ret_value(self): 582 | a = """\ 583 | def nop() -> int: 584 | pass 585 | """ 586 | self.unchanged(a) 587 | 588 | def test_idempotency_arg1_ret_value(self): 589 | a = """\ 590 | def nop(a) -> int: 591 | pass 592 | """ 593 | self.unchanged(a) 594 | 595 | def test_idempotency_arg1_default_1(self): 596 | a = """\ 597 | def nop(a: int=0): 598 | pass 599 | """ 600 | self.unchanged(a) 601 | 602 | def test_idempotency_arg1_default_2(self): 603 | a = """\ 604 | def nop(a: List[int]=[]): 605 | pass 606 | """ 607 | self.unchanged(a) 608 | 609 | def test_idempotency_arg1_default_3(self): 610 | a = """\ 611 | def nop(a: List[int]=[1,2,3]): 612 | pass 613 | """ 614 | self.unchanged(a) 615 | 616 | 617 | 618 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | mypy_extensions>=0.3.0 2 | pytest>=3.3.0 3 | setuptools>=28.8.0 4 | six>=1.11.0 5 | typing>=3.6.2; python_version < '3.5' 6 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal = 1 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os 4 | from setuptools import setup 5 | 6 | def get_long_description(): 7 | filename = os.path.join(os.path.dirname(__file__), 'README.md') 8 | with open(filename) as f: 9 | return f.read() 10 | 11 | setup(name='pyannotate', 12 | version='1.2.0', 13 | description="PyAnnotate: Auto-generate PEP-484 annotations", 14 | long_description=get_long_description(), 15 | long_description_content_type="text/markdown", 16 | author='Dropbox', 17 | author_email='guido@dropbox.com', 18 | url='https://github.com/dropbox/pyannotate', 19 | license='Apache 2.0', 20 | platforms=['POSIX'], 21 | packages=['pyannotate_runtime', 'pyannotate_tools', 22 | 'pyannotate_tools.annotations', 'pyannotate_tools.fixes'], 23 | entry_points={'console_scripts': ['pyannotate=pyannotate_tools.annotations.__main__:main']}, 24 | classifiers=[ 25 | 'Development Status :: 3 - Alpha', 26 | 'Environment :: Console', 27 | 'Intended Audience :: Developers', 28 | 'License :: OSI Approved :: Apache Software License', 29 | 'Operating System :: POSIX', 30 | 'Programming Language :: Python :: 2', 31 | 'Programming Language :: Python :: 2.7', 32 | 'Programming Language :: Python :: 3', 33 | 'Programming Language :: Python :: 3.4', 34 | 'Programming Language :: Python :: 3.5', 35 | 'Programming Language :: Python :: 3.6', 36 | 'Programming Language :: Python :: 3.7', 37 | 'Topic :: Software Development', 38 | ], 39 | install_requires = ['six', 40 | 'mypy_extensions', 41 | 'typing >= 3.5.3; python_version < "3.5"' 42 | ], 43 | ) 44 | -------------------------------------------------------------------------------- /tests/integration_test.py: -------------------------------------------------------------------------------- 1 | """Some things you just can't test as unit tests""" 2 | 3 | import os 4 | import subprocess 5 | import sys 6 | import tempfile 7 | import unittest 8 | import shutil 9 | 10 | 11 | example = """ 12 | def main(): 13 | print(gcd(15, 10)) 14 | print(gcd(45, 12)) 15 | 16 | def gcd(a, b): 17 | while b: 18 | a, b = b, a%b 19 | return a 20 | """ 21 | 22 | driver = """ 23 | from pyannotate_runtime import collect_types 24 | 25 | if __name__ == '__main__': 26 | collect_types.init_types_collection() 27 | with collect_types.collect(): 28 | main() 29 | collect_types.dump_stats('type_info.json') 30 | """ 31 | 32 | class_example = """ 33 | class A(object): pass 34 | 35 | def f(x): 36 | return x 37 | 38 | def main(): 39 | f(A()) 40 | f(A()) 41 | """ 42 | 43 | 44 | class IntegrationTest(unittest.TestCase): 45 | 46 | def setUp(self): 47 | self.savedir = os.getcwd() 48 | os.putenv('PYTHONPATH', self.savedir) 49 | self.tempdir = tempfile.mkdtemp() 50 | os.chdir(self.tempdir) 51 | 52 | def tearDown(self): 53 | os.chdir(self.savedir) 54 | shutil.rmtree(self.tempdir) 55 | 56 | def test_simple(self): 57 | with open('gcd.py', 'w') as f: 58 | f.write(example) 59 | with open('driver.py', 'w') as f: 60 | f.write('from gcd import main\n') 61 | f.write(driver) 62 | subprocess.check_call([sys.executable, 'driver.py']) 63 | output = subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', 'gcd.py']) 64 | lines = output.splitlines() 65 | assert b'+ # type: () -> None' in lines 66 | assert b'+ # type: (int, int) -> int' in lines 67 | 68 | def test_auto_any(self): 69 | with open('gcd.py', 'w') as f: 70 | f.write(example) 71 | output = subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', '-a', 'gcd.py']) 72 | lines = output.splitlines() 73 | assert b'+ # type: () -> None' in lines 74 | assert b'+ # type: (Any, Any) -> Any' in lines 75 | 76 | def test_no_type_info(self): 77 | with open('gcd.py', 'w') as f: 78 | f.write(example) 79 | try: 80 | subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', 'gcd.py'], 81 | stderr=subprocess.STDOUT) 82 | assert False, "Expected an error" 83 | except subprocess.CalledProcessError as err: 84 | assert err.returncode == 1 85 | lines = err.output.splitlines() 86 | assert (b"Can't open type info file: " 87 | b"[Errno 2] No such file or directory: 'type_info.json'" in lines) 88 | 89 | def test_package(self): 90 | os.makedirs('foo') 91 | with open('foo/__init__.py', 'w') as f: 92 | pass 93 | with open('foo/gcd.py', 'w') as f: 94 | f.write(example) 95 | with open('driver.py', 'w') as f: 96 | f.write('from foo.gcd import main\n') 97 | f.write(driver) 98 | subprocess.check_call([sys.executable, 'driver.py']) 99 | output = subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', 'foo/gcd.py']) 100 | lines = output.splitlines() 101 | assert b'+ # type: () -> None' in lines 102 | assert b'+ # type: (int, int) -> int' in lines 103 | 104 | def test_subdir(self): 105 | os.makedirs('foo') 106 | with open('foo/gcd.py', 'w') as f: 107 | f.write(example) 108 | with open('driver.py', 'w') as f: 109 | f.write('import sys\n') 110 | f.write('sys.path.insert(0, "foo")\n') 111 | f.write('from gcd import main\n') 112 | f.write(driver) 113 | subprocess.check_call([sys.executable, 'driver.py']) 114 | output = subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', 115 | # Construct platform-correct pathname: 116 | os.path.join('foo', 'gcd.py')]) 117 | lines = output.splitlines() 118 | assert b'+ # type: () -> None' in lines 119 | assert b'+ # type: (int, int) -> int' in lines 120 | 121 | def test_subdir_w_class(self): 122 | os.makedirs('foo') 123 | with open('foo/bar.py', 'w') as f: 124 | f.write(class_example) 125 | with open('driver.py', 'w') as f: 126 | f.write('import sys\n') 127 | f.write('sys.path.insert(0, "foo")\n') 128 | f.write('from bar import main\n') 129 | f.write(driver) 130 | subprocess.check_call([sys.executable, 'driver.py']) 131 | output = subprocess.check_output([sys.executable, '-m', 'pyannotate_tools.annotations', 132 | # Construct platform-correct pathname: 133 | os.path.join('foo', 'bar.py')]) 134 | lines = output.splitlines() 135 | print(b'\n'.join(lines).decode()) 136 | assert b'+ # type: () -> None' in lines 137 | assert b'+ # type: (A) -> A' in lines 138 | assert not any(line.startswith(b'+') and b'import' in line for line in lines) 139 | --------------------------------------------------------------------------------