Initial commit

.gitignore

@@ -0,0 +1,118 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+Session.vim
+.vim/
+.cache/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Viliam Valent
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,45 @@
+# KeyOf
+
+[Mypy](https://mypy-lang.org/) plugin for static type checking of [TypedDict](https://peps.python.org/pep-0589/) keys
+inspired by TypeScript's [`keyof` type operator](https://www.typescriptlang.org/docs/handbook/2/keyof-types.html).
+
+## Installation
+
+```shell
+pip install keyof
+```
+
+### Mypy plugin
+
+Add `keyof.mypy_plugin` to the list of plugins in your [mypy config file](https://mypy.readthedocs.io/en/latest/config_file.html)
+(for example `pyproject.toml`)
+
+```toml
+[tool.mypy]
+plugins = ["keyof.mypy_plugin"]
+```
+
+## Usage
+
+```python
+from typing import TypedDict
+
+from keyof import KeyOf
+
+
+class Data(TypedDict):
+    version: int
+    command: str
+
+
+def get_data(data: Data, key: KeyOf[Data]) -> int | str:
+    return data[key]
+
+
+data = Data(version=1, command="foo")
+
+get_data(data, "version")  # OK
+
+get_data(data, "not_existing_key") # error
+# Argument 2 to "get_data" has incompatible type "Literal['not_existing_key']"; expected "Literal['version', 'command']"
+```

keyof/__init__.py

@@ -0,0 +1,11 @@
+from typing import Any, Generic, NoReturn, TypeVar
+
+_T = TypeVar("_T")
+
+
+class KeyOf(Any, Generic[_T]):
+    def __new__(cls, *_args: Any, **_kwargs: Any) -> NoReturn:
+        raise TypeError("Type KeyOf cannot be instantiated.")
+
+    def __init_subclass__(cls, *_args: Any, **_kwargs: Any) -> NoReturn:
+        raise TypeError("Cannot subclass KeyOf")

keyof/compat.py

@@ -0,0 +1,3 @@
+from typing import Any, TypeAlias
+
+KeyOf: TypeAlias = Any

keyof/mypy_plugin.py

@@ -0,0 +1,65 @@
+from collections.abc import Callable, Sequence
+from typing import Final
+
+from mypy.errorcodes import TYPE_ARG
+from mypy.plugin import AnalyzeTypeContext, Plugin, TypeAnalyzerPluginInterface
+from mypy.types import (
+    LiteralType,
+    Type,
+    TypeAliasType,
+    TypedDictType,
+    UninhabitedType,
+    UnionType,
+)
+
+_MIN_MYPY_VERSION = (1, 0, 0)
+_TYPE_NAME: Final = "KeyOf"
+_TYPES_TO_ANALYZE: Final = frozenset(["keyof.KeyOf", "keyof.compat.KeyOf"])
+
+
+def _create_string_literal(value: str, api: TypeAnalyzerPluginInterface) -> LiteralType:
+    return LiteralType(value, api.named_type("builtins.str", []))
+
+
+def _create_literal_type(values: Sequence[str], api: TypeAnalyzerPluginInterface) -> UnionType | LiteralType:
+    if len(values) == 1:
+        return _create_string_literal(values[0], api)
+    return UnionType([_create_string_literal(value, api) for value in values])
+
+
+def _analyze_typed_dict_key(ctx: AnalyzeTypeContext) -> Type:
+    if (args_len := len(ctx.type.args)) != 1:
+        ctx.api.fail(
+            f'"{_TYPE_NAME}" expects 1 type argument, but {args_len} given',
+            ctx.context,
+            code=TYPE_ARG,
+        )
+        return UninhabitedType()
+    argument = ctx.type.args[0]
+    analyzed: Type | None = ctx.api.analyze_type(argument)
+    if isinstance(analyzed, TypeAliasType):
+        analyzed = analyzed.expand_all_if_possible()
+    if not isinstance(analyzed, TypedDictType):
+        ctx.api.fail(
+            f'Argument 1 to "{_TYPE_NAME}" has incompatible type "{analyzed}"; expected "TypedDict"',
+            ctx.context,
+        )
+        return UninhabitedType()
+    return _create_literal_type(list(analyzed.items.keys()), ctx.api)
+
+
+class CustomPlugin(Plugin):
+    def get_type_analyze_hook(self, fullname: str) -> Callable[[AnalyzeTypeContext], Type] | None:
+        if fullname in _TYPES_TO_ANALYZE:
+            return _analyze_typed_dict_key
+        return None
+
+
+def _version_str_to_tuple(version: str) -> tuple[int, ...]:
+    return tuple(int(part) for part in version.partition("+")[0].split("."))
+
+
+def plugin(version: str) -> type[Plugin]:
+    if _version_str_to_tuple(version) >= _MIN_MYPY_VERSION:
+        return CustomPlugin
+    return Plugin