├── compytype
    ├── __init__.py
    └── plugin.py
├── .gitignore
├── mypy.ini
├── tests
    └── test_add.py
└── README.md


/compytype/__init__.py:
--------------------------------------------------------------------------------
1 | 


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | __pycache__
2 | .pytest_cache
3 | .mypy_cache
4 | 


--------------------------------------------------------------------------------
/mypy.ini:
--------------------------------------------------------------------------------
1 | # Global options:
2 | 
3 | [mypy]
4 | plugins = compytype/plugin.py
5 | ignore_missing_imports = True
6 | 


--------------------------------------------------------------------------------
/tests/test_add.py:
--------------------------------------------------------------------------------
1 | import pytest
2 | from typing import Literal
3 | 
4 | @pytest.mark.mypy_testing
5 | def test_annotated_add():
6 |     op1: Literal[3] = 3
7 |     op2: Literal[4] = 4
8 |     reveal_type(op1 + op2) # N: Revealed type is 'Literal[7]'
9 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # ComPyType
 2 | 
 3 | Implementation of our "Type-level Computations for Ruby Libraries" in Python. Python programmers regularly use dataframes, matrices etc. and type-level computation can be a major benefit for catching bugs early.
 4 | 
 5 | This uses the mypy type checker and builds on top of it as a plugin.
 6 | 
 7 | ### Test
 8 | 
 9 | ```
10 | pytest
11 | ```
12 | 


--------------------------------------------------------------------------------
/compytype/plugin.py:
--------------------------------------------------------------------------------
 1 | from mypy.types import LiteralType
 2 | from mypy.plugin import Plugin
 3 | 
 4 | def type_add(ctx, *args):
 5 |     trecv = ctx.type
 6 |     targs = ctx.arg_types
 7 |     if trecv.__class__ is LiteralType and len(targs) == 1 and targs[0][0].__class__ is LiteralType:
 8 |         ret = trecv.value + targs[0][0].value
 9 |         return LiteralType(ret, fallback=ctx.api.named_generic_type('builtins.int', []))
10 |     else:
11 |         return None
12 | 
13 | TYPE_DICT = {
14 |     'builtins.int.__add__': type_add
15 | }
16 | 
17 | class ComPyTypePlugin(Plugin):
18 | 
19 |     def get_method_hook(self, fullname):
20 |         return TYPE_DICT.get(fullname, None)
21 | 
22 | def plugin(version: str):
23 |     # ignore version argument if the plugin works with all mypy versions.
24 |     return ComPyTypePlugin
25 | 


--------------------------------------------------------------------------------