├── MANIFEST.in ├── setup.py ├── tox.ini ├── pyproject.toml ├── sorcery ├── __init__.py ├── core.py └── spells.py ├── make_release.sh ├── LICENSE.txt ├── .github └── workflows │ └── pytest.yml ├── setup.cfg ├── .gitignore ├── README.md └── tests.py /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include LICENSE.txt 2 | include README.md 3 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | if __name__ == "__main__": 4 | setup() 5 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py35,py36,py37,py38,py39,py310,pypy36 3 | 4 | [testenv] 5 | commands = python tests.py 6 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools", "wheel", "setuptools_scm[toml]"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [tool.setuptools_scm] 6 | write_to = "sorcery/version.py" 7 | write_to_template = "__version__ = '{version}'" 8 | -------------------------------------------------------------------------------- /sorcery/__init__.py: -------------------------------------------------------------------------------- 1 | from .core import spell, no_spells 2 | 3 | from .spells import unpack_keys, unpack_attrs, args_with_source, dict_of, print_args, call_with_name, delegate_to_attr, \ 4 | maybe, select_from, magic_kwargs, assigned_names, switch, timeit 5 | 6 | try: 7 | from .version import __version__ 8 | except ImportError: # pragma: no cover 9 | # version.py is auto-generated with the git tag when building 10 | __version__ = "???" 11 | -------------------------------------------------------------------------------- /make_release.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -eux 3 | 4 | # Ensure that there are no uncommitted changes 5 | # which would mess up using the git tag as a version 6 | [ -z "$(git status --porcelain)" ] 7 | 8 | if [ -z "${1+x}" ] 9 | then 10 | set +x 11 | echo Provide a version argument 12 | echo "${0} .." 13 | exit 1 14 | else 15 | if [[ ${1} =~ ^([0-9]+)(\.[0-9]+)?(\.[0-9]+)?$ ]]; then 16 | : 17 | else 18 | echo "Not a valid release tag." 19 | exit 1 20 | fi 21 | fi 22 | 23 | tox -p auto 24 | 25 | export TAG="v${1}" 26 | git tag "${TAG}" 27 | git push origin master "${TAG}" 28 | rm -rf ./build ./dist 29 | python -m build --sdist --wheel . 30 | twine upload ./dist/*.whl dist/*.tar.gz 31 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2021 Alex Hall 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /.github/workflows/pytest.yml: -------------------------------------------------------------------------------- 1 | name: Tests 2 | on: [push, pull_request] 3 | jobs: 4 | build: 5 | runs-on: ubuntu-latest 6 | strategy: 7 | matrix: 8 | python-version: [3.5, 3.6, 3.7, 3.8, 3.9-dev, 3.10-dev, 'pypy-3.6'] 9 | steps: 10 | - uses: actions/checkout@v2 11 | - name: Set up Python ${{ matrix.python-version }} 12 | uses: actions/setup-python@v2 13 | with: 14 | python-version: ${{ matrix.python-version }} 15 | - name: run tests 16 | run: | 17 | pip install --upgrade pip 18 | pip install --upgrade setuptools setuptools_scm pep517 coveralls 19 | pip install . 20 | coverage run --source sorcery tests.py 21 | coverage report -m 22 | - name: Coveralls Python 23 | uses: AndreMiras/coveralls-python-action@v20201129 24 | with: 25 | parallel: true 26 | flag-name: test-${{ matrix.python-version }} 27 | coveralls_finish: 28 | needs: build 29 | runs-on: ubuntu-latest 30 | steps: 31 | - name: Coveralls Finished 32 | uses: AndreMiras/coveralls-python-action@v20201129 33 | with: 34 | parallel-finished: true 35 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | name = sorcery 3 | author = Alex Hall 4 | author_email = alex.mojaki@gmail.com 5 | license = MIT 6 | description = Dark magic delights in Python 7 | long_description = file: README.md 8 | long_description_content_type = text/markdown 9 | url = https://github.com/alexmojaki/sorcery 10 | classifiers = 11 | License :: OSI Approved :: MIT License 12 | Programming Language :: Python 13 | Programming Language :: Python :: 3 14 | Programming Language :: Python :: 3.5 15 | Programming Language :: Python :: 3.6 16 | Programming Language :: Python :: 3.7 17 | Programming Language :: Python :: 3.8 18 | Programming Language :: Python :: 3.9 19 | Programming Language :: Python :: 3.10 20 | Operating System :: OS Independent 21 | Programming Language :: Python :: Implementation :: CPython 22 | Programming Language :: Python :: Implementation :: PyPy 23 | 24 | [options] 25 | packages = sorcery 26 | include_package_data = True 27 | setup_requires = setuptools; setuptools_scm[toml] 28 | install_requires = 29 | executing 30 | littleutils>=0.2.1 31 | asttokens 32 | wrapt 33 | 34 | [coverage:run] 35 | relative_files = True 36 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Created by .ignore support plugin (hsz.mobi) 2 | ### Python template 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | .hypothesis/ 50 | 51 | # Translations 52 | *.mo 53 | *.pot 54 | 55 | # Django stuff: 56 | *.log 57 | .static_storage/ 58 | .media/ 59 | local_settings.py 60 | 61 | # Flask stuff: 62 | instance/ 63 | .webassets-cache 64 | 65 | # Scrapy stuff: 66 | .scrapy 67 | 68 | # Sphinx documentation 69 | docs/_build/ 70 | 71 | # PyBuilder 72 | target/ 73 | 74 | # Jupyter Notebook 75 | .ipynb_checkpoints 76 | 77 | # pyenv 78 | .python-version 79 | 80 | # celery beat schedule file 81 | celerybeat-schedule 82 | 83 | # SageMath parsed files 84 | *.sage.py 85 | 86 | # Environments 87 | .env 88 | .venv 89 | env/ 90 | venv/ 91 | ENV/ 92 | env.bak/ 93 | venv.bak/ 94 | 95 | # Spyder project settings 96 | .spyderproject 97 | .spyproject 98 | 99 | # Rope project settings 100 | .ropeproject 101 | 102 | # mkdocs documentation 103 | /site 104 | 105 | # mypy 106 | .mypy_cache/ 107 | 108 | /sorcery/version.py 109 | -------------------------------------------------------------------------------- /sorcery/core.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import sys 3 | from functools import lru_cache, partial 4 | from typing import Tuple 5 | 6 | from executing import only, Source 7 | 8 | NAMED_EXPR_SUPPORT = sys.version_info.major == 3 and sys.version_info.minor >= 8 9 | 10 | class FrameInfo(object): 11 | """ 12 | Contains metadata about where a spell is being called. 13 | An instance of this is passed as the first argument to any spell. 14 | Users should not instantiate this class themselves. 15 | 16 | There are two essential attributes: 17 | 18 | - frame: the execution frame in which the spell is being called 19 | - call: the ast.Call node where the spell is being called 20 | See https://greentreesnakes.readthedocs.io/en/latest/nodes.html 21 | to learn how to navigate the AST 22 | """ 23 | 24 | def __init__(self, executing): 25 | self.frame = executing.frame 26 | self.executing = executing 27 | self.call = executing.node 28 | 29 | def assigned_names(self, *, 30 | allow_one: bool = False, 31 | allow_loops: bool = False 32 | ) -> Tuple[Tuple[str], ast.AST]: 33 | """ 34 | Calls the function assigned_names for this instance's Call node. 35 | """ 36 | return assigned_names(self.call, 37 | allow_one=allow_one, 38 | allow_loops=allow_loops) 39 | 40 | def get_source(self, node: ast.AST) -> str: 41 | """ 42 | Returns a string containing the source code of an AST node in the 43 | same file as this call. 44 | """ 45 | return self.executing.source.asttokens().get_text(node) 46 | 47 | 48 | @lru_cache() 49 | def statement_containing_node(node: ast.AST) -> ast.stmt: 50 | while not isinstance(node, ast.stmt): 51 | node = node.parent 52 | return node 53 | 54 | 55 | @lru_cache() 56 | def assigned_names(node, *, 57 | allow_one: bool, 58 | allow_loops: bool 59 | ) -> Tuple[Tuple[str], ast.AST]: 60 | """ 61 | Finds the names being assigned to in the nearest ancestor of 62 | the given node that assigns names and satisfies the given conditions. 63 | 64 | If allow_loops is false, this only considers assignment statements, 65 | e.g. `x, y = ...`. If it's true, then for loops and comprehensions are 66 | also considered. 67 | 68 | If allow_one is false, nodes which assign only one name are ignored. 69 | 70 | Returns: 71 | 1. a tuple of strings containing the names of the nodes being assigned 72 | 2. The AST node where the assignment happens 73 | """ 74 | 75 | while hasattr(node, 'parent'): 76 | node = node.parent 77 | 78 | target = None 79 | 80 | if NAMED_EXPR_SUPPORT and isinstance(node, ast.NamedExpr): 81 | target = node.target 82 | elif isinstance(node, ast.Assign): 83 | target = only(node.targets) 84 | elif isinstance(node, (ast.For, ast.comprehension)) and allow_loops: 85 | target = node.target 86 | 87 | if not target: 88 | continue 89 | 90 | names = node_names(target) 91 | if len(names) > 1 or allow_one: 92 | break 93 | else: 94 | raise TypeError('No assignment found') 95 | 96 | return names, node 97 | 98 | 99 | def node_names(node: ast.AST) -> Tuple[str]: 100 | """ 101 | Returns a tuple of strings containing the names of 102 | the nodes under the given node. 103 | 104 | The node must be a tuple or list literal, or a single named node. 105 | 106 | See the doc of the function node_name. 107 | """ 108 | if isinstance(node, (ast.Tuple, ast.List)): 109 | names = tuple(node_name(x) for x in node.elts) 110 | else: 111 | names = (node_name(node),) 112 | return names 113 | 114 | 115 | def node_name(node: ast.AST) -> str: 116 | """ 117 | Returns the 'name' of a node, which is either: 118 | - the name of a variable 119 | - the name of an attribute 120 | - the contents of a string literal key, as in d['key'] 121 | """ 122 | if isinstance(node, ast.Name): 123 | return node.id 124 | elif isinstance(node, ast.Attribute): 125 | return node.attr 126 | elif isinstance(node, ast.Subscript): 127 | index = node.slice 128 | if isinstance(index, ast.Index): 129 | index = index.value 130 | if isinstance(index, ast.Str): 131 | return index.s 132 | raise TypeError('Cannot extract name from %s' % node) 133 | 134 | 135 | class Spell(object): 136 | """ 137 | A Spell is a special callable that has information about where it's being 138 | called from. 139 | 140 | To create a spell, decorate a function with @spell. 141 | An instance of FrameInfo will be passed to the first argument of the function, 142 | while the other arguments will come from the call. For example: 143 | 144 | @spell 145 | def my_spell(frame_info, foo): 146 | ... 147 | 148 | will be called as just `my_spell(foo)`. 149 | """ 150 | 151 | _excluded_codes = set() 152 | 153 | # Called when decorating a function 154 | def __init__(self, func): 155 | self.func = func 156 | 157 | # Called when a spell is accessed as an attribute 158 | # (see the descriptor protocol) 159 | def __get__(self, instance, owner): 160 | # Functions are descriptors, which allow methods to 161 | # automatically bind the self argument. 162 | # Here we have to manually invoke that. 163 | method = self.func.__get__(instance, owner) 164 | return Spell(method) 165 | 166 | def at(self, frame_info: FrameInfo): 167 | """ 168 | Returns a callable that has frame_info already bound as the first argument 169 | of the spell, and will accept only the other arguments normally. 170 | 171 | Use this to use one spell inside another. 172 | """ 173 | return partial(self.func, frame_info) 174 | 175 | # Called when the spell is called 'plainly', e.g. my_spell(foo), 176 | # i.e. just as a variable without being an attribute of anything. 177 | # Calls where the spell is an attribute go throuh __get__. 178 | def __call__(self, *args, **kwargs): 179 | frame = sys._getframe(1) 180 | 181 | while frame.f_code in self._excluded_codes: 182 | frame = frame.f_back 183 | 184 | executing = Source.executing(frame) 185 | assert executing.node, "Failed to find call node" 186 | return self.at(FrameInfo(executing))(*args, **kwargs) 187 | 188 | def __repr__(self): 189 | return '%s(%r)' % ( 190 | self.__class__.__name__, 191 | self.func 192 | ) 193 | 194 | 195 | spell = Spell 196 | 197 | 198 | def no_spells(func): 199 | """ 200 | Decorate a function with this to indicate that no spells are used 201 | directly in this function, but the function may be used to 202 | access a spell dynamically. Spells looking for where they are being 203 | called from will skip the decorated function and look at the enclosing 204 | frame instead. 205 | 206 | For example, suppose you have a class with a method that is a spell, 207 | e.g.: 208 | 209 | class A: 210 | @ magic_kwargs # makes foo a spell 211 | def foo(self, **kwargs): 212 | pass 213 | 214 | And another class that wraps the first: 215 | 216 | class B: 217 | def __init__(self): 218 | self.a = A() 219 | 220 | And you want users to call foo without going through A, you could write: 221 | 222 | @no_spells 223 | def foo(self, **kwargs): 224 | self.a.foo(**kwargs) 225 | 226 | Note that the method B.foo must have the same name (foo) as the spell A.foo. 227 | 228 | Or, if you wanted to delegate all unknown attributes to `self.a`, you could write: 229 | 230 | @no_spells 231 | def __getattr__(self, item): 232 | return getattr(self.a, item) 233 | 234 | In either case `B().foo(...)` will work as expected. 235 | """ 236 | 237 | Spell._excluded_codes.add(func.__code__) 238 | return func 239 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # sorcery 2 | 3 | [![Build Status](https://github.com/alexmojaki/sorcery/actions/workflows/pytest.yml/badge.svg)](https://github.com/alexmojaki/sorcery/actions/workflows/pytest.yml) [![Coverage Status](https://coveralls.io/repos/github/alexmojaki/sorcery/badge.svg?branch=master)](https://coveralls.io/github/alexmojaki/sorcery?branch=master) [![Supports Python 3.5+, including PyPy](https://img.shields.io/pypi/pyversions/sorcery.svg)](https://pypi.python.org/pypi/sorcery) 4 | 5 | This package lets you use and write callables called 'spells' that know where they're being called from and can use that information to do otherwise impossible things. 6 | 7 | Note: previously spells had a complicated implementation that placed limitations on how they could be called. Now spells are just a thin wrapper around [`executing`](https://github.com/alexmojaki/executing) which is much better. You may be better off using `executing` directly depending on your use case. This repo is now mostly just a fun collection of things to do with it. 8 | 9 | * [Installation](#installation) 10 | * [Quick examples](#quick-examples) 11 | * [`assigned_names`](#assigned_names) 12 | * [`unpack_keys` and `unpack_attrs`](#unpack_keys-and-unpack_attrs) 13 | * [`dict_of`](#dict_of) 14 | * [`print_args`](#print_args) 15 | * [`call_with_name` and `delegate_to_attr`](#call_with_name-and-delegate_to_attr) 16 | * [`maybe`](#maybe) 17 | * [`timeit`](#timeit) 18 | * [`switch`](#switch) 19 | * [`select_from`](#select_from) 20 | * [How to write your own spells](#how-to-write-your-own-spells) 21 | * [Using other spells within spells](#using-other-spells-within-spells) 22 | * [Other helpers](#other-helpers) 23 | * [Should I actually use this library?](#should-i-actually-use-this-library) 24 | 25 | ## Installation 26 | 27 | pip install sorcery 28 | 29 | ## Quick examples 30 | 31 | See the docstrings for more detail. 32 | 33 | from sorcery import (assigned_names, unpack_keys, unpack_attrs, 34 | dict_of, print_args, call_with_name, 35 | delegate_to_attr, maybe, select_from) 36 | 37 | ### `assigned_names` 38 | 39 | Instead of: 40 | 41 | foo = func('foo') 42 | bar = func('bar') 43 | 44 | write: 45 | 46 | foo, bar = [func(name) for name in assigned_names()] 47 | 48 | Instead of: 49 | 50 | class Thing(Enum): 51 | foo = 'foo' 52 | bar = 'bar' 53 | 54 | write: 55 | 56 | class Thing(Enum): 57 | foo, bar = assigned_names() 58 | 59 | ### `unpack_keys` and `unpack_attrs` 60 | 61 | Instead of: 62 | 63 | foo = d['foo'] 64 | bar = d['bar'] 65 | 66 | write: 67 | 68 | foo, bar = unpack_keys(d) 69 | 70 | Similarly, instead of: 71 | 72 | foo = x.foo 73 | bar = x.bar 74 | 75 | write: 76 | 77 | foo, bar = unpack_attrs(x) 78 | 79 | ### `dict_of` 80 | 81 | Instead of: 82 | 83 | dict(foo=foo, bar=bar, spam=thing()) 84 | 85 | write: 86 | 87 | dict_of(foo, bar, spam=thing()) 88 | 89 | (see also: `magic_kwargs`) 90 | 91 | ### `print_args` 92 | 93 | For easy debugging, instead of: 94 | 95 | print("foo =", foo) 96 | print("bar() =", bar()) 97 | 98 | write: 99 | 100 | print_args(foo, bar()) 101 | 102 | To write your own version of this (e.g. if you want to add colour), use `args_with_source`. 103 | 104 | If you like this, I recommend the `pp` function in the [`snoop`](https://github.com/alexmojaki/snoop) library. 105 | 106 | ### `call_with_name` and `delegate_to_attr` 107 | 108 | Sometimes you want to create many similar methods which differ only in a string argument which is equal to the name of the method. Given this class: 109 | 110 | ```python 111 | class C: 112 | def generic(self, method_name, *args, **kwargs): 113 | ... 114 | ``` 115 | 116 | Inside the class definition, instead of: 117 | 118 | ```python 119 | def foo(self, x, y): 120 | return self.generic('foo', x, y) 121 | 122 | def bar(self, z): 123 | return self.generic('bar', z) 124 | ``` 125 | 126 | write: 127 | 128 | ```python 129 | foo, bar = call_with_name(generic) 130 | ``` 131 | 132 | For a specific common use case: 133 | 134 | ```python 135 | class Wrapper: 136 | def __init__(self, thing): 137 | self.thing = thing 138 | 139 | def foo(self, x, y): 140 | return self.thing.foo(x, y) 141 | 142 | def bar(self, z): 143 | return self.thing.bar(z) 144 | ``` 145 | 146 | you can instead write: 147 | 148 | ```python 149 | foo, bar = delegate_to_attr('thing') 150 | ``` 151 | 152 | For a more concrete example, here is a class that wraps a list and has all the usual list methods while ensuring that any methods which usually create a new list actually create a new wrapper: 153 | 154 | ```python 155 | class MyListWrapper(object): 156 | def __init__(self, lst): 157 | self.list = lst 158 | 159 | def _make_new_wrapper(self, method_name, *args, **kwargs): 160 | method = getattr(self.list, method_name) 161 | new_list = method(*args, **kwargs) 162 | return type(self)(new_list) 163 | 164 | append, extend, clear, __repr__, __str__, __eq__, __hash__, \ 165 | __contains__, __len__, remove, insert, pop, index, count, \ 166 | sort, __iter__, reverse, __iadd__ = spells.delegate_to_attr('list') 167 | 168 | copy, __add__, __radd__, __mul__, __rmul__ = spells.call_with_name(_make_new_wrapper) 169 | ``` 170 | 171 | Of course, there are less magical DRY ways to accomplish this (e.g. looping over some strings and using `setattr`), but they will not tell your IDE/linter what methods `MyListWrapper` has or doesn't have. 172 | 173 | ### `maybe` 174 | 175 | While we wait for the `?.` operator from [PEP 505](https://www.python.org/dev/peps/pep-0505/), here's an alternative. Instead of: 176 | 177 | None if foo is None else foo.bar() 178 | 179 | write: 180 | 181 | maybe(foo).bar() 182 | 183 | If you want a slightly less magical version, consider [pymaybe](https://github.com/ekampf/pymaybe). 184 | 185 | ### `timeit` 186 | 187 | Instead of 188 | 189 | ```python 190 | import timeit 191 | 192 | nums = [3, 1, 2] 193 | setup = 'from __main__ import nums' 194 | 195 | print(timeit.repeat('min(nums)', setup)) 196 | print(timeit.repeat('sorted(nums)[0]', setup)) 197 | ``` 198 | 199 | write: 200 | 201 | ```python 202 | import sorcery 203 | 204 | nums = [3, 1, 2] 205 | 206 | if sorcery.timeit(): 207 | result = min(nums) 208 | else: 209 | result = sorted(nums)[0] 210 | ``` 211 | 212 | ### `switch` 213 | 214 | Instead of: 215 | 216 | ```python 217 | if val == 1: 218 | x = 1 219 | elif val == 2 or val == bar(): 220 | x = spam() 221 | elif val == dangerous_function(): 222 | x = spam() * 2 223 | else: 224 | x = -1 225 | ``` 226 | 227 | write: 228 | 229 | ```python 230 | x = switch(val, lambda: { 231 | 1: 1, 232 | {{ 2, bar() }}: spam(), 233 | dangerous_function(): spam() * 2 234 | }, default=-1) 235 | ``` 236 | 237 | This really will behave like the if/elif chain above. The dictionary is just 238 | some nice syntax, but no dictionary is ever actually created. The keys 239 | are evaluated only as needed, in order, and only the matching value is evaluated. 240 | 241 | ### `select_from` 242 | 243 | Instead of: 244 | 245 | ```python 246 | cursor.execute(''' 247 | SELECT foo, bar 248 | FROM my_table 249 | WHERE spam = ? 250 | AND thing = ? 251 | ''', [spam, thing]) 252 | 253 | for foo, bar in cursor: 254 | ... 255 | ``` 256 | 257 | write: 258 | 259 | ```python 260 | for foo, bar in select_from('my_table', where=[spam, thing]): 261 | ... 262 | ``` 263 | 264 | ## How to write your own spells 265 | 266 | Decorate a function with `@spell`. An instance of the class `FrameInfo` will be passed to the first argument of the function, while the other arguments will come from the call. For example: 267 | 268 | ```python 269 | from sorcery import spell 270 | 271 | @spell 272 | def my_spell(frame_info, foo): 273 | ... 274 | ``` 275 | 276 | will be called as just `my_spell(foo)`. 277 | 278 | The most important piece of information you are likely to use is `frame_info.call`. This is the `ast.Call` node where the spell is being called. [Here](https://greentreesnakes.readthedocs.io/en/latest/nodes.html) is some helpful documentation for navigating the AST. Every node also has a `parent` attribute added to it. 279 | 280 | `frame_info.frame` is the execution frame in which the spell is being called - see the [inspect](https://docs.python.org/3/library/inspect.html) docs for what you can do with this. 281 | 282 | Those are the essentials. See [the source](https://github.com/alexmojaki/sorcery/blob/master/sorcery/spells.py) of various spells for some examples, it's not that complicated. 283 | 284 | ### Using other spells within spells 285 | 286 | Sometimes you want to reuse the magic of one spell in another spell. Simply calling the other spell won't do what you want - you want to tell the other spell to act as if it's being called from the place your own spell is called. For this, add insert `.at(frame_info)` between the spell you're using and its arguments. 287 | 288 | Let's look at a concrete example. Here's the definition of the spell `args_with_source`: 289 | 290 | ```python 291 | @spell 292 | def args_with_source(frame_info, *args): 293 | """ 294 | Returns a list of pairs of: 295 | - the source code of the argument 296 | - the value of the argument 297 | for each argument. 298 | 299 | For example: 300 | 301 | args_with_source(foo(), 1+2) 302 | 303 | is the same as: 304 | 305 | [ 306 | ("foo()", foo()), 307 | ("1+2", 3) 308 | ] 309 | """ 310 | ... 311 | ``` 312 | 313 | The magic of `args_with_source` is that it looks at its arguments wherever it's called and extracts their source code. Here is a simplified implementation of the `print_args` spell which uses that magic: 314 | 315 | ```python 316 | @spell 317 | def simple_print_args(frame_info, *args): 318 | for source, arg in args_with_source.at(frame_info)(*args): 319 | print(source, '=', arg) 320 | ``` 321 | 322 | Then when you call `simple_print_args(foo(), 1+2)`, the `Call` node of that expression will be passed down to `args_with_source.at(frame_info)` so that the source is extracted from the correct arguments. Simply writing `args_with_source(*args)` would be wrong, as that would give the source `"*args"`. 323 | 324 | ### Other helpers 325 | 326 | That's all you really need to get started writing a spell, but here are pointers to some other stuff that might help. See the docstrings for details. 327 | 328 | The module `sorcery.core` has these helper functions: 329 | 330 | - `node_names(node: ast.AST) -> Tuple[str]` 331 | - `node_name(node: ast.AST) -> str` 332 | - `statement_containing_node(node: ast.AST) -> ast.stmt:` 333 | 334 | `FrameInfo` has these methods: 335 | 336 | - `assigned_names(...)` 337 | - `get_source(self, node: ast.AST) -> str` 338 | 339 | ## Should I actually use this library? 340 | 341 | If you're still getting the hang of Python, no. This will lead to confusion about what is normal and expected in Python and will hamper your learning. 342 | 343 | In a serious business or production context, I wouldn't recommend most of the spells unless you're quite careful. Their unusual nature may confuse other readers of the code, and tying the behaviour of your code to things like the names of variables may not be good for readability and refactoring. There are some exceptions though: 344 | 345 | - `call_with_name` and `delegate_to_attr` 346 | - `assigned_names` for making `Enum`s. 347 | - `print_args` when debugging 348 | 349 | If you're writing code where performance and stability aren't critical, e.g. if it's for fun or you just want to get some code down as fast as possible and you can polish it later, then go for it. 350 | 351 | The point of this library is not just to be used in actual code. It's a way to explore and think about API and language design, readability, and the limits of Python itself. It was fun to create and I hope others can have fun playing around with it. Come [have a chat](https://gitter.im/python-sorcery/Lobby) about what spells you think would be cool, what features you wish Python had, or what crazy projects you want to create. 352 | 353 | If you're interested in this stuff, particularly creative uses of the Python AST, you may also be interested in: 354 | 355 | - [executing](https://github.com/alexmojaki/executing) the backbone of this library 356 | - [snoop](https://github.com/alexmojaki/snoop): a feature-rich and convenient debugging library which also uses `executing` as well as various other magic and tricks 357 | - [birdseye](https://github.com/alexmojaki/birdseye): a debugger which records the value of every expression 358 | - [MacroPy](https://github.com/lihaoyi/macropy): syntactic macros in Python by transforming the AST at import time 359 | -------------------------------------------------------------------------------- /tests.py: -------------------------------------------------------------------------------- 1 | import ast 2 | import sqlite3 3 | import sys 4 | import traceback 5 | import unittest 6 | from io import StringIO 7 | from time import sleep 8 | from unittest import mock 9 | 10 | from littleutils import SimpleNamespace, only 11 | 12 | import sorcery as spells 13 | from sorcery import unpack_keys, unpack_attrs, print_args, magic_kwargs, maybe, args_with_source, spell 14 | from sorcery.spells import PYPY 15 | 16 | 17 | class MyListWrapper(object): 18 | def __init__(self, lst): 19 | self.list = lst 20 | 21 | def _make_new_wrapper(self, method_name, *args, **kwargs): 22 | method = getattr(self.list, method_name) 23 | new_list = method(*args, **kwargs) 24 | return type(self)(new_list) 25 | 26 | append, extend, clear, __repr__, __str__, __eq__, __hash__, \ 27 | __contains__, __len__, remove, insert, pop, index, count, \ 28 | sort, __iter__, reverse, __iadd__ = spells.delegate_to_attr('list') 29 | 30 | copy, __add__, __radd__, __mul__, __rmul__ = spells.call_with_name(_make_new_wrapper) 31 | 32 | 33 | class Foo(object): 34 | @magic_kwargs 35 | def bar(self, **kwargs): 36 | return set(kwargs.items()) | {self} 37 | 38 | 39 | @magic_kwargs 40 | def magic_only_kwarg(n, *, y): 41 | return n, y 42 | 43 | 44 | class TestStuff(unittest.TestCase): 45 | def test_unpack_keys_basic(self): 46 | obj = SimpleNamespace(thing=SimpleNamespace()) 47 | d = dict(foo=1, bar=3, spam=7, baz=8, x=9) 48 | out = {} 49 | foo, obj.thing.spam, obj.bar, out['baz'] = unpack_keys(d) 50 | self.assertEqual(foo, d['foo']) 51 | self.assertEqual(obj.bar, d['bar']) 52 | self.assertEqual(obj.thing.spam, d['spam']) 53 | self.assertEqual(out, {'baz': d['baz']}) 54 | 55 | def test_unpack_keys_for_loop(self): 56 | results = [] 57 | for x, y in unpack_keys([ 58 | dict(x=1, y=2), 59 | dict(x=3, z=4), 60 | dict(a=5, y=6), 61 | dict(b=7, c=8), 62 | ], default=999): 63 | results.append((x, y)) 64 | self.assertEqual(results, [ 65 | (1, 2), 66 | (3, 999), 67 | (999, 6), 68 | (999, 999), 69 | ]) 70 | 71 | def test_unpack_keys_list_comprehension(self): 72 | self.assertEqual( 73 | [(y, x) for x, y in unpack_keys([ 74 | dict(x=1, y=2), 75 | dict(x=3, y=4), 76 | ])], 77 | [ 78 | (2, 1), 79 | (4, 3), 80 | ]) 81 | 82 | def test_unpack_keys_bigger_expression(self): 83 | x, y = map(int, unpack_keys(dict(x='1', y='2'))) 84 | self.assertEqual(x, 1) 85 | self.assertEqual(y, 2) 86 | 87 | def test_unpack_keys_skip_single_assigned_name(self): 88 | x, y = [int(v) for v in unpack_keys(dict(x='1', y='2'))] 89 | self.assertEqual(x, 1) 90 | self.assertEqual(y, 2) 91 | 92 | def test_unpack_keys_extras(self): 93 | env = dict(DATABASE_USERNAME='me', 94 | DATABASE_PASSWORD='secret') 95 | username, password = unpack_keys(env, prefix='DATABASE_', swapcase=True) 96 | self.assertEqual(username, 'me') 97 | self.assertEqual(password, 'secret') 98 | 99 | def test_unpack_attrs(self): 100 | obj = SimpleNamespace(aa='bv', bb='cc', cc='aa') 101 | cc, bb, aa = unpack_attrs(obj) 102 | self.assertEqual(aa, obj.aa) 103 | self.assertEqual(bb, obj.bb) 104 | self.assertEqual(cc, obj.cc) 105 | 106 | d, e = unpack_attrs(obj, default=9) 107 | assert d == e == 9 108 | 109 | def test_print_args(self): 110 | out = StringIO() 111 | x = 3 112 | y = 4 113 | print_args(x + y, 114 | x * y, 115 | x - 116 | y, file=out) 117 | self.assertEqual('''\ 118 | x + y = 119 | 7 120 | 121 | x * y = 122 | 12 123 | 124 | x - 125 | y = 126 | -1 127 | 128 | ''', out.getvalue()) 129 | 130 | def test_dict_of(self): 131 | a = 1 132 | obj = SimpleNamespace(b=2) 133 | self.assertEqual(spells.dict_of( 134 | a, obj.b, 135 | c=3, d=4 136 | ), dict( 137 | a=a, b=obj.b, 138 | c=3, d=4)) 139 | 140 | def test_no_starargs_in_dict_of(self): 141 | args = [1, 2] 142 | with self.assertRaises(TypeError): 143 | spells.dict_of(*args) 144 | 145 | def test_delegation(self): 146 | lst = MyListWrapper([1, 2, 3]) 147 | lst.append(4) 148 | lst.extend([1, 2]) 149 | lst = (lst + [5]).copy() 150 | self.assertEqual(type(lst), MyListWrapper) 151 | self.assertEqual(lst, [1, 2, 3, 4, 1, 2, 5]) 152 | 153 | def test_magic_kwargs(self): 154 | foo = Foo() 155 | x = 1 156 | y = 2 157 | w = 10 158 | self.assertEqual(foo.bar(x, y, z=3), 159 | {('x', x), ('y', y), ('z', 3), foo}) 160 | 161 | self.assertEqual(magic_only_kwarg(x, y), (x, y)) 162 | 163 | @magic_kwargs 164 | def spam(n, **kwargs): 165 | return n, kwargs 166 | 167 | self.assertEqual(spam(x, y, z=5), 168 | (x, dict(y=y, z=5))) 169 | 170 | @magic_kwargs 171 | def spam(n, m, **kwargs): 172 | return n, m, kwargs 173 | 174 | self.assertEqual(spam(x, w, y, z=5), 175 | (x, w, dict(y=y, z=5))) 176 | 177 | with self.assertRaises(TypeError): 178 | @magic_kwargs 179 | def _(a=1): 180 | print(a) 181 | 182 | with self.assertRaises(TypeError): 183 | @magic_kwargs 184 | def _(*a): 185 | print(a) 186 | 187 | def test_maybe(self): 188 | if PYPY: 189 | with self.assertRaises(NotImplementedError): 190 | maybe(None) 191 | return 192 | 193 | n = None 194 | assert maybe(n) is None 195 | self.assertIsNone(maybe(n)) 196 | assert maybe(n).a.b.c()[4]().asd.asd()() is None 197 | assert maybe(n)()()() is None 198 | assert maybe(0) == 0 199 | assert maybe({'a': 3})['a'] == 3 200 | assert maybe({'a': {'b': 3}})['a']['b'] == 3 201 | assert maybe({'a': {'b': 3}})['a']['b'] + 2 == 5 202 | assert maybe({'a': {'b': None}})['a']['b'] is None 203 | 204 | def test_select_from(self): 205 | conn = sqlite3.connect(':memory:') 206 | c = conn.cursor() 207 | c.execute('CREATE TABLE points (x INT, y INT)') 208 | c.execute("INSERT INTO points VALUES (5, 3), (8, 1)") 209 | conn.commit() 210 | 211 | assert [(3, 5), (1, 8)] == [(y, x) for y, x in spells.select_from('points')] 212 | y = 1 213 | x = spells.select_from('points', where=[y]) 214 | assert (x, y) == (8, 1) 215 | 216 | def test_multiple_attr_calls(self): 217 | x = 3 218 | y = 5 219 | self.assertEqual([ 220 | spells.dict_of(x), 221 | spells.dict_of(y), 222 | ], [dict(x=x), dict(y=y)]) 223 | 224 | self.assertEqual([spells.dict_of(x), spells.dict_of(y)], 225 | [dict(x=x), dict(y=y)]) 226 | 227 | def test_no_assignment(self): 228 | with self.assertRaises(TypeError): 229 | unpack_keys(dict(x=1, y=2)) 230 | 231 | def test_spell_repr(self): 232 | self.assertRegex(repr(spells.dict_of), 233 | r'Spell\(\)') 234 | 235 | def test_assigned_names(self): 236 | x, y = ['_' + s for s in spells.assigned_names()] 237 | self.assertEqual(x, '_x') 238 | self.assertEqual(y, '_y') 239 | 240 | # noinspection PyTrailingSemicolon 241 | def test_semicolons(self): 242 | # @formatter:off 243 | tester(1); tester(2); tester(3) 244 | tester(9 245 | ); tester( 246 | 8); tester( 247 | 99 248 | ); tester(33); tester([4, 249 | 5, 6, [ 250 | 7]]) 251 | # @formatter:on 252 | 253 | def test_args_with_source(self): 254 | self.assertEqual(args_with_source(1 + 2, 3 * 4), 255 | [("1 + 2", 3), ("3 * 4", 12)]) 256 | self.assertEqual( 257 | args_with_source( 258 | self.assertEqual(args_with_source(1 + 2), [("1 + 2", 3)])), 259 | [( 260 | 'self.assertEqual(args_with_source(1 + 2), [("1 + 2", 3)])', 261 | None, 262 | )], 263 | ) 264 | 265 | def test_switch(self): 266 | result = spells.switch(2, lambda: { 267 | 1: 10, 268 | 2: 20, 269 | 1 / 0: 1 / 0 270 | }) 271 | self.assertEqual(result, 20) 272 | 273 | result = spells.switch(2, lambda: { 274 | 1: 10, 275 | {{5, 2, 1 / 0}}: 20, 276 | 3: 1 / 0 277 | }) 278 | self.assertEqual(result, 20) 279 | 280 | with self.assertRaises(KeyError): 281 | spells.switch(2, lambda: { 282 | 1: 10, 283 | 3: 30, 284 | }) 285 | 286 | result = spells.switch(2, lambda: { 287 | 1: 10, 288 | 3: 30, 289 | }, default=-1) 290 | self.assertEqual(result, -1) 291 | 292 | with self.assertRaises(TypeError): 293 | spells.switch(2, { 294 | 1: 10, 295 | 2: 20, 296 | }) 297 | 298 | with self.assertRaises(TypeError): 299 | spells.switch(2, lambda: [{ 300 | 1: 10, 301 | 2: 20, 302 | }]) 303 | 304 | def test_timeit_in_function(self): 305 | with self.assertRaises(ValueError): 306 | spells.timeit() 307 | 308 | def test_decorator(self): 309 | @empty_decorator 310 | @decorator_with_args(tester('123'), x=int()) 311 | @tester(list(tuple([1, 2])), returns=empty_decorator) 312 | @tester( 313 | list( 314 | tuple( 315 | [3, 4])), 316 | returns=empty_decorator) 317 | @empty_decorator 318 | @decorator_with_args( 319 | str(), 320 | x=int()) 321 | @tester(list(tuple([5, 6])), returns=empty_decorator) 322 | @tester(list(tuple([7, 8])), returns=empty_decorator) 323 | @empty_decorator 324 | @decorator_with_args(tester('sdf'), x=tester('123234')) 325 | def foo(): 326 | pass 327 | 328 | def test_list_comprehension(self): 329 | str([tester(int(x)) for x in tester([1]) for _ in tester([2]) for __ in [3]]) 330 | str([[[tester(int(x)) for x in tester([1])] for _ in tester([2])] for __ in [3]]) 331 | return str([(1, [ 332 | (2, [ 333 | tester(int(x)) for x in tester([1])]) 334 | for _ in tester([2])]) 335 | for __ in [3]]) 336 | 337 | def test_lambda(self): 338 | self.assertEqual((lambda x: (tester(x), tester(x)))(tester(3)), (3, 3)) 339 | (lambda: (lambda: tester(1))())() 340 | self.assertEqual((lambda: [tester(x) for x in tester([1, 2])])(), [1, 2]) 341 | 342 | def test_indirect_call(self): 343 | dict(x=tester)['x'](tester)(3) 344 | 345 | def test_compound_statements(self): 346 | with self.assertRaises(TypeError): 347 | try: 348 | for _ in tester([2]): 349 | while tester(0): 350 | pass 351 | else: 352 | tester(4) 353 | else: 354 | tester(5) 355 | raise ValueError 356 | except tester(ValueError): 357 | tester(9) 358 | raise TypeError 359 | finally: 360 | tester(10) 361 | 362 | # PyCharm getting confused somehow? 363 | # noinspection PyUnreachableCode 364 | str() 365 | 366 | with self.assertRaises(tester(Exception)): 367 | if tester(0): 368 | pass 369 | elif tester(0): 370 | pass 371 | elif tester(1 / 0): 372 | pass 373 | 374 | def test_generator(self): 375 | def gen(): 376 | for x in [1, 2]: 377 | yield tester(x) 378 | 379 | gen2 = (tester(x) for x in tester([1, 2])) 380 | 381 | assert list(gen()) == list(gen2) == [1, 2] 382 | 383 | 384 | @spell 385 | def tester(frame_info, arg, returns=None): 386 | result = eval( 387 | compile(ast.Expression(only(frame_info.call.args)), '<>', 'eval'), 388 | frame_info.frame.f_globals, 389 | frame_info.frame.f_locals, 390 | ) 391 | assert result == arg, (result, arg) 392 | if returns is None: 393 | return arg 394 | return returns 395 | 396 | 397 | assert tester([1, 2, 3]) == [1, 2, 3] 398 | 399 | 400 | def empty_decorator(f): 401 | return f 402 | 403 | 404 | def decorator_with_args(*_, **__): 405 | return empty_decorator 406 | 407 | 408 | class TestTimeit(unittest.TestCase): 409 | def patch(self, *args, **kwargs): 410 | patcher = mock.patch(*args, **kwargs) 411 | patcher.start() 412 | self.addCleanup(patcher.stop) 413 | 414 | def setUp(self): 415 | self.patch('sorcery.spells._raise', lambda e: e) 416 | self.patch('sys.stdout', StringIO()) 417 | 418 | def assert_usual_output(self): 419 | self.assertRegex( 420 | sys.stdout.getvalue(), 421 | r""" 422 | Number of trials: 1 423 | 424 | Method 1: 1\.\d{3} 425 | Method 2: 1\.\d{3} 426 | 427 | Method 1: 1\.\d{3} 428 | Method 2: 1\.\d{3} 429 | 430 | Best times: 431 | ----------- 432 | Method 1: 1\.\d{3} 433 | Method 2: 1\.\d{3} 434 | """.strip()) 435 | 436 | def test_no_result(self): 437 | if spells.timeit(repeat=2): 438 | sleep(1) 439 | else: 440 | sleep(1.1) 441 | self.assert_usual_output() 442 | 443 | # noinspection PyUnusedLocal 444 | def test_matching_result(self): 445 | if spells.timeit(repeat=2): 446 | sleep(1) 447 | result = 3 448 | else: 449 | sleep(1.1) 450 | result = 3 451 | self.assert_usual_output() 452 | 453 | # noinspection PyUnusedLocal 454 | def test_not_matching_result(self): 455 | with self.assertRaises(AssertionError): 456 | if spells.timeit(): 457 | result = 3 458 | else: 459 | result = 4 460 | 461 | def test_exception(self): 462 | try: 463 | if spells.timeit(): 464 | print(1 / 0) 465 | else: 466 | pass 467 | except ZeroDivisionError: 468 | traceback.print_exc(file=sys.stdout) 469 | 470 | stdout = sys.stdout.getvalue() 471 | self.assertIn('', stdout) 472 | self.assertIn('1 / 0', stdout) 473 | 474 | 475 | if __name__ == '__main__': 476 | unittest.main() 477 | -------------------------------------------------------------------------------- /sorcery/spells.py: -------------------------------------------------------------------------------- 1 | from __future__ import generator_stop 2 | 3 | import ast 4 | import operator 5 | import sys 6 | import timeit as real_timeit 7 | import unittest 8 | from functools import lru_cache 9 | from inspect import signature 10 | from io import StringIO 11 | from itertools import chain 12 | from pprint import pprint 13 | from textwrap import dedent 14 | 15 | import wrapt 16 | from littleutils import only 17 | from sorcery.core import spell, node_names, node_name 18 | 19 | _NO_DEFAULT = object() 20 | PYPY = 'pypy' in sys.version.lower() 21 | 22 | 23 | @spell 24 | def assigned_names(frame_info): 25 | """ 26 | Instead of: 27 | 28 | foo = func('foo') 29 | bar = func('bar') 30 | 31 | write: 32 | 33 | foo, bar = map(func, assigned_names()) 34 | 35 | or: 36 | 37 | foo, bar = [func(name) for name in assigned_names()] 38 | 39 | Instead of: 40 | 41 | class Thing(Enum): 42 | foo = 'foo' 43 | bar = 'bar' 44 | 45 | write: 46 | 47 | class Thing(Enum): 48 | foo, bar = assigned_names() 49 | 50 | More generally, this function returns a tuple of strings representing the names being assigned to. 51 | 52 | The result can be assigned to any combination of either: 53 | 54 | - plain variables, 55 | - attributes, or 56 | - subscripts (square bracket access) with string literal keys 57 | 58 | So the following: 59 | 60 | spam, x.foo, y['bar'] = assigned_names() 61 | 62 | is equivalent to: 63 | 64 | spam = 'spam' 65 | x.foo = 'foo' 66 | y['bar'] = 'bar' 67 | 68 | Any expression is allowed to the left of the attribute/subscript. 69 | 70 | Only simple tuple unpacking is allowed: 71 | 72 | - no nesting, e.g. (a, b), c = ... 73 | - no stars, e.g. a, *b = ... 74 | - no chains, e.g. a, b = c = ... 75 | - no assignment to a single name without unpacking, e.g. a = ... 76 | """ 77 | return frame_info.assigned_names()[0] 78 | 79 | 80 | @spell 81 | def unpack_keys(frame_info, x, default=_NO_DEFAULT, prefix=None, swapcase=False): 82 | """ 83 | Instead of: 84 | 85 | foo = d['foo'] 86 | bar = d['bar'] 87 | 88 | write: 89 | 90 | foo, bar = unpack_keys(d) 91 | 92 | Instead of: 93 | 94 | foo = d.get('foo', 0) 95 | bar = d.get('bar', 0) 96 | 97 | write: 98 | 99 | foo, bar = unpack_keys(d, default=0) 100 | 101 | Instead of: 102 | 103 | foo = d['data_foo'] 104 | bar = d['data_bar'] 105 | 106 | write: 107 | 108 | foo, bar = unpack_keys(d, prefix='data_') 109 | 110 | Instead of: 111 | 112 | foo = d['FOO'] 113 | bar = d['BAR'] 114 | 115 | write: 116 | 117 | foo, bar = unpack_keys(d, swapcase=True) 118 | 119 | and similarly, instead of: 120 | 121 | FOO = d['foo'] 122 | BAR = d['bar'] 123 | 124 | write: 125 | 126 | FOO, BAR = unpack_keys(d, swapcase=True) 127 | 128 | Note that swapcase is not applied to the prefix, so for example you should write: 129 | 130 | env = dict(DATABASE_USERNAME='me', 131 | DATABASE_PASSWORD='secret') 132 | username, password = unpack_keys(env, prefix='DATABASE_', swapcase=True) 133 | 134 | The rules of the assigned_names spell apply. 135 | 136 | This can be seamlessly used in for loops, even inside comprehensions, e.g. 137 | 138 | for foo, bar in unpack_keys(list_of_dicts): 139 | ... 140 | 141 | If there are multiple assignment targets in the statement, e.g. if you have 142 | a nested list comprehension, the target nearest to the function call will 143 | determine the keys. For example, the keys 'foo' and 'bar' will be extracted in: 144 | 145 | [[foo + bar + y for foo, bar in unpack_keys(x)] 146 | for x, y in z] 147 | 148 | Like assigned_names, the unpack call can be part of a bigger expression, 149 | and the assignment will still be found. So for example instead of: 150 | 151 | foo = int(d['foo']) 152 | bar = int(d['bar']) 153 | 154 | you can write: 155 | 156 | foo, bar = map(int, unpack_keys(d)) 157 | 158 | or: 159 | 160 | foo, bar = [int(v) for v in unpack_keys(d)] 161 | 162 | The second version works because the spell looks for multiple names being assigned to, 163 | so it doesn't just unpack 'v'. 164 | 165 | """ 166 | 167 | if default is _NO_DEFAULT: 168 | getter = operator.getitem 169 | else: 170 | # Essentially dict.get, without relying on that method existing 171 | def getter(d, name): 172 | try: 173 | return d[name] 174 | except KeyError: 175 | return default 176 | 177 | return _unpack(frame_info, x, getter, prefix, swapcase) 178 | 179 | 180 | @spell 181 | def unpack_attrs(frame_info, x, default=_NO_DEFAULT, prefix=None, swapcase=False): 182 | """ 183 | This is similar to unpack_keys, but for attributes. 184 | 185 | Instead of: 186 | 187 | foo = x.foo 188 | bar = x.bar 189 | 190 | write: 191 | 192 | foo, bar = unpack_attrs(x) 193 | """ 194 | 195 | if default is _NO_DEFAULT: 196 | getter = getattr 197 | else: 198 | def getter(d, name): 199 | return getattr(d, name, default) 200 | 201 | return _unpack(frame_info, x, getter, prefix, swapcase) 202 | 203 | 204 | def _unpack(frame_info, x, getter, prefix, swapcase): 205 | names, node = frame_info.assigned_names(allow_loops=True) 206 | 207 | def fix_name(n): 208 | if swapcase: 209 | n = n.swapcase() 210 | if prefix: 211 | n = prefix + n 212 | return n 213 | 214 | if isinstance(node, ast.Assign): 215 | return [getter(x, fix_name(name)) for name in names] 216 | else: # for loop 217 | return ([getter(d, fix_name(name)) for name in names] 218 | for d in x) 219 | 220 | 221 | @spell 222 | def args_with_source(frame_info, *args): 223 | """ 224 | Returns a list of pairs of: 225 | - the source code of the argument 226 | - the value of the argument 227 | for each argument. 228 | 229 | For example: 230 | 231 | args_with_source(foo(), 1+2) 232 | 233 | is the same as: 234 | 235 | [ 236 | ("foo()", foo()), 237 | ("1+2", 3) 238 | ] 239 | """ 240 | return [ 241 | (frame_info.get_source(arg), value) 242 | for arg, value in zip(frame_info.call.args, args) 243 | ] 244 | 245 | 246 | @spell 247 | def dict_of(frame_info, *args, **kwargs): 248 | """ 249 | Instead of: 250 | 251 | {'foo': foo, 'bar': bar, 'spam': thing()} 252 | 253 | or: 254 | 255 | dict(foo=foo, bar=bar, spam=thing()) 256 | 257 | write: 258 | 259 | dict_of(foo, bar, spam=thing()) 260 | 261 | In other words, returns a dictionary with an item for each argument, 262 | where positional arguments use their names as keys, 263 | and keyword arguments do the same as in the usual dict constructor. 264 | 265 | The positional arguments can be any of: 266 | 267 | - plain variables, 268 | - attributes, or 269 | - subscripts (square bracket access) with string literal keys 270 | 271 | So the following: 272 | 273 | dict_of(spam, x.foo, y['bar']) 274 | 275 | is equivalent to: 276 | 277 | dict(spam=spam, foo=x.foo, bar=y['bar']) 278 | 279 | *args are not allowed. 280 | 281 | To give your own functions the ability to turn positional argments into 282 | keyword arguments, use the decorator magic_kwargs. 283 | 284 | """ 285 | 286 | result = { 287 | node_name(arg): value 288 | for arg, value in zip(frame_info.call.args[-len(args):], args) 289 | } 290 | result.update(kwargs) 291 | return result 292 | 293 | 294 | @spell 295 | def print_args(frame_info, *args, file=None): 296 | """ 297 | For each argument, prints the source code of that argument 298 | and its value. Returns the first argument. 299 | """ 300 | for source, arg in args_with_source.at(frame_info)(*args): 301 | print(source + ' =', file=file) 302 | pprint(arg, stream=file) 303 | print(file=file) 304 | return args and args[0] 305 | 306 | 307 | @spell 308 | def call_with_name(frame_info, func): 309 | """ 310 | Given: 311 | 312 | class C: 313 | def generic(self, method_name, *args, **kwargs): 314 | ... 315 | 316 | Inside the class definition, instead of: 317 | 318 | def foo(self, x, y): 319 | return self.generic('foo', x, y) 320 | 321 | def bar(self, z): 322 | return self.generic('bar', z) 323 | 324 | write: 325 | 326 | foo, bar = call_with_name(generic) 327 | 328 | This only works for methods inside classes, not free functions. 329 | """ 330 | def make_func(name): 331 | return lambda self, *args, **kwargs: func(self, name, *args, **kwargs) 332 | 333 | return [ 334 | make_func(name) 335 | for name in frame_info.assigned_names()[0] 336 | ] 337 | 338 | 339 | @spell 340 | def delegate_to_attr(frame_info, attr_name): 341 | """ 342 | This is a special case of the use case fulfilled by call_with_name. 343 | 344 | Given: 345 | 346 | class Wrapper: 347 | def __init__(self, thing): 348 | self.thing = thing 349 | 350 | Inside the class definition, instead of: 351 | 352 | def foo(self, x, y): 353 | return self.thing.foo(x, y) 354 | 355 | def bar(self, z): 356 | return self.thing.bar(z) 357 | 358 | Write: 359 | 360 | foo, bar = delegate_to_attr('thing') 361 | 362 | Specifically, this will make: 363 | 364 | Wrapper().foo 365 | 366 | equivalent to: 367 | 368 | Wrapper().thing.foo 369 | """ 370 | def make_func(name): 371 | return property(lambda self: getattr(getattr(self, attr_name), name)) 372 | 373 | return [ 374 | make_func(name) 375 | for name in frame_info.assigned_names()[0] 376 | ] 377 | 378 | 379 | class _Nothing(object): 380 | def __init__(self, count): 381 | self.__count = count 382 | 383 | def __getattribute__(self, item): 384 | if item == '_Nothing__count': 385 | return object.__getattribute__(self, item) 386 | return _Nothing.__op(self) 387 | 388 | def __op(self, *_args, **_kwargs): 389 | self.__count -= 1 390 | if self.__count == 0: 391 | return None 392 | 393 | return self 394 | 395 | __getitem__ = __call__ = __op 396 | 397 | 398 | @spell 399 | def maybe(frame_info, x): 400 | """ 401 | Instead of: 402 | 403 | None if foo is None else foo.bar() 404 | 405 | write: 406 | 407 | maybe(foo).bar() 408 | 409 | Specifically, if foo is not None, then maybe(foo) is just foo. 410 | 411 | If foo is None, then any sequence of attributes, subscripts, or 412 | calls immediately to the right of maybe(foo) is ignored, and 413 | the final result is None. So maybe(foo)[0].x.y.bar() is None, 414 | while func(maybe(foo)[0].x.y.bar()) is func(None) because enclosing 415 | expressions are not affected. 416 | """ 417 | if x is not None: 418 | return x 419 | 420 | node = frame_info.call 421 | count = 0 422 | while True: 423 | parent = node.parent 424 | if not (isinstance(parent, ast.Attribute) or 425 | isinstance(parent, ast.Call) and parent.func is node or 426 | isinstance(parent, ast.Subscript) and parent.value is node): 427 | break 428 | count += 1 429 | node = parent 430 | 431 | if count == 0: 432 | return x 433 | 434 | return _Nothing(count) 435 | 436 | 437 | if PYPY: 438 | def maybe(_): 439 | raise NotImplementedError("maybe isn't supported on pypy`") 440 | 441 | 442 | @spell 443 | def select_from(frame_info, sql, params=(), cursor=None, where=None): 444 | """ 445 | Instead of: 446 | 447 | cursor.execute(''' 448 | SELECT foo, bar 449 | FROM my_table 450 | WHERE spam = ? 451 | AND thing = ? 452 | ''', [spam, thing]) 453 | 454 | for foo, bar in cursor: 455 | ... 456 | 457 | write: 458 | 459 | for foo, bar in select_from('my_table', where=[spam, thing]): 460 | ... 461 | 462 | Specifically: 463 | - the assigned names (similar to the assigned_names and unpack_keys spells) 464 | are placed in the SELECT clause 465 | - the first argument (usually just a table name but can be any SQL) 466 | goes after the FROM 467 | - if the where argument is supplied, it must be a list or tuple literal of values 468 | which are supplied as query parameters and whose names are used in a 469 | WHERE clause using the = and AND operators. 470 | If you use this argument, don't put a WHERE clause in the sql argument and 471 | don't supply params 472 | - a cursor object is automatically pulled from the calling frame, but if this 473 | doesn't work you can supply one with the cursor keyword argument 474 | - the params argument can be supplied for more custom cases than the where 475 | argument provides. 476 | - if this is used in a loop or list comprehension, all rows in the result 477 | will be iterated over. 478 | If it is used in an assignment statement, one row will be returned. 479 | - If there are multiple names being assigned (i.e. multiple columns being selected) 480 | then the row will be returned and thus unpacked. If there is only one name, 481 | it will automatically be unpacked so you don't have to add [0]. 482 | 483 | This spell is much more a fun rough idea than the others. It is expected that there 484 | are many use cases it will not fit into nicely. 485 | """ 486 | if cursor is None: 487 | frame = frame_info.frame 488 | cursor = only(c for c in chain(frame.f_locals.values(), 489 | frame.f_globals.values()) 490 | if 'cursor' in str(type(c).__mro__).lower() and 491 | callable(getattr(c, 'execute', None))) 492 | names, node = frame_info.assigned_names(allow_one=True, allow_loops=True) 493 | sql = 'SELECT %s FROM %s' % (', '.join(names), sql) 494 | 495 | if where: 496 | where_arg = only(kw.value for kw in frame_info.call.keywords 497 | if kw.arg == 'where') 498 | where_names = node_names(where_arg) 499 | assert len(where_names) == len(where) 500 | sql += ' WHERE ' + ' AND '.join('%s = ?' % name for name in where_names) 501 | assert params == () 502 | params = where 503 | 504 | cursor.execute(sql, params) 505 | 506 | def unpack(row): 507 | if len(row) == 1: 508 | return row[0] 509 | else: 510 | return row 511 | 512 | if isinstance(node, ast.Assign): 513 | return unpack(cursor.fetchone()) 514 | else: 515 | def vals(): 516 | for row in cursor: 517 | yield unpack(row) 518 | 519 | return vals() 520 | 521 | 522 | def magic_kwargs(func): 523 | """ 524 | Applying this decorator allows a function to interpret positional 525 | arguments as keyword arguments, using the name of the positional argument 526 | as the keyword. For example, given: 527 | 528 | @magic_kwargs 529 | def func(*, foo, bar, spam): 530 | 531 | or 532 | 533 | @magic_kwargs 534 | def func(**kwargs): 535 | 536 | then instead of: 537 | 538 | func(foo=foo, bar=bar, spam=thing) 539 | 540 | you can just write: 541 | 542 | func(foo, bar, spam=thing) 543 | 544 | Without the @magic_kwargs, the closest magical alternative would be: 545 | 546 | func(**dict_of(foo, bar, spam=thing)) 547 | 548 | The function is not allowed to have optional positional parameters, e.g. 549 | `def func(x=1)`, or *args. 550 | """ 551 | 552 | args_count = 0 553 | for param in signature(func).parameters.values(): 554 | if (param.kind == param.VAR_POSITIONAL or 555 | param.kind == param.POSITIONAL_OR_KEYWORD and 556 | param.default != param.empty): 557 | raise TypeError( 558 | 'The type of the parameter %s is not allowed with @magic_kwargs' 559 | % param.name) 560 | if param.kind == param.POSITIONAL_OR_KEYWORD: 561 | args_count += 1 562 | 563 | @wrapt.decorator 564 | def wrapper(wrapped, instance, args, kwargs): 565 | frame_info, *args = args 566 | count = args_count - (instance is not None) # account for self argument 567 | normal_args = args[:count] 568 | magic_args = args[count:] 569 | full_kwargs = dict_of.at(frame_info)(*magic_args, **kwargs) 570 | return wrapped(*normal_args, **full_kwargs) 571 | 572 | return spell(wrapper(func)) 573 | 574 | 575 | @spell 576 | def switch(frame_info, val, _cases, *, default=_NO_DEFAULT): 577 | """ 578 | Instead of: 579 | 580 | if val == 1: 581 | x = 1 582 | elif val == 2 or val == bar(): 583 | x = spam() 584 | elif val == dangerous_function(): 585 | x = spam() * 2 586 | else: 587 | x = -1 588 | 589 | write: 590 | 591 | x = switch(val, lambda: { 592 | 1: 1, 593 | {{ 2, bar() }}: spam(), 594 | dangerous_function(): spam() * 2 595 | }, default=-1) 596 | 597 | This really will behave like the if/elif chain above. The dictionary is just 598 | some nice syntax, but no dictionary is ever actually created. The keys 599 | are evaluated only as needed, in order, and only the matching value is evaluated. 600 | The keys are not hashed, only compared for equality, so non-hashable keys like lists 601 | are allowed. 602 | 603 | If the default is not specified and no matching value is found, a KeyError is raised. 604 | 605 | Note that `if val == 2 or val == bar()` is translated to `{{ 2, bar() }}`. 606 | This is to allow emulating multiple case clauses for the same block as in 607 | the switch construct in other languages. The double braces {{}} create a value 608 | that's impossible to evaluate normally (a set containing a set) so that it's clear 609 | we don't simply want to check `val == {{ 2, bar() }}`, whereas `{2, bar()}` would be 610 | evaluated and checked normally. 611 | As always, the contents are lazily evaluated and compared in order. 612 | 613 | The keys and values are evaluated with the compiler statement 614 | `from __future__ import generator_stop` in effect (which you should really be 615 | considering using anyway if you're using Python < 3.7). 616 | 617 | """ 618 | 619 | frame = frame_info.frame 620 | switcher = _switcher(frame_info.call.args[1], frame.f_code) 621 | 622 | def ev(k): 623 | return eval(k, frame.f_globals, frame.f_locals) 624 | 625 | def check(k): 626 | return ev(k) == val 627 | 628 | for key_code, value_code in switcher: 629 | if isinstance(key_code, tuple): 630 | test = any(map(check, key_code)) 631 | else: 632 | test = check(key_code) 633 | if test: 634 | return ev(value_code) 635 | 636 | if default is _NO_DEFAULT: 637 | raise KeyError(val) 638 | else: 639 | return default 640 | 641 | 642 | @lru_cache() 643 | def _switcher(cases, f_code): 644 | if not (isinstance(cases, ast.Lambda) and 645 | isinstance(cases.body, ast.Dict)): 646 | raise TypeError('The second argument to switch must be a lambda with no arguments ' 647 | 'that returns a dictionary literal') 648 | 649 | def comp(node): 650 | return compile(ast.Expression(node), 651 | filename=f_code.co_filename, 652 | mode='eval') 653 | 654 | result = [] 655 | for key, value in zip(cases.body.keys, 656 | cases.body.values): 657 | 658 | if (isinstance(key, ast.Set) and 659 | isinstance(key.elts[0], ast.Set)): 660 | key_code = tuple(comp(k) for k in key.elts[0].elts) 661 | else: 662 | key_code = comp(key) 663 | 664 | result.append((key_code, comp(value))) 665 | return result 666 | 667 | 668 | def _raise(e): 669 | # for tests 670 | raise e 671 | 672 | 673 | class TimerWithExc(real_timeit.Timer): 674 | def timeit(self, *args, **kwargs): 675 | try: 676 | return super().timeit(*args, **kwargs) 677 | except: 678 | # Sets up linecache for future tracebacks 679 | self.print_exc(StringIO()) 680 | raise 681 | 682 | 683 | @spell 684 | def timeit(frame_info, repeat=5): 685 | """ 686 | This function is for writing quick scripts for comparing the speeds 687 | of two snippets of code that do the same thing. It's a nicer interface 688 | to the standard timeit module that doesn't require putting your code in strings, so you can 689 | use your IDE features, while still using the standard timeit for accuracy. 690 | 691 | Instead of 692 | 693 | import timeit 694 | 695 | nums = [3, 1, 2] 696 | setup = 'from __main__ import nums' 697 | 698 | print(timeit.repeat('min(nums)', setup)) 699 | print(timeit.repeat('sorted(nums)[0]', setup)) 700 | 701 | write: 702 | 703 | import sorcery 704 | 705 | nums = [3, 1, 2] 706 | 707 | if sorcery.timeit(): 708 | result = min(nums) 709 | else: 710 | result = sorted(nums)[0] 711 | 712 | The if statement is just syntax for denoting the two blocks of code 713 | being tested. Some other nice features of this function over the standard 714 | timeit: 715 | 716 | - Automatically determines a high enough 'number' argument. 717 | - Asserts that any variable named 'result' is equal in both snippets, 718 | for correctness testing. The variable should be present in both or 719 | neither snippets. 720 | - Nice formatting of results for easy comparison, including best times 721 | - Source lines shown in tracebacks 722 | 723 | The spell must be called at the top level of a module, not inside 724 | another function definition. 725 | """ 726 | 727 | globs = frame_info.frame.f_globals 728 | if globs is not frame_info.frame.f_locals: 729 | _raise(ValueError('Must execute in global scope')) 730 | 731 | setup = 'from %s import %s\n' % ( 732 | globs['__name__'], 733 | ', '.join(globs.keys()), 734 | ) 735 | if_stmt = frame_info.call.parent 736 | stmts = [ 737 | dedent('\n'.join(map(frame_info.get_source, lines))) 738 | for lines in [if_stmt.body, if_stmt.orelse] 739 | ] 740 | 741 | timers = [ 742 | TimerWithExc(stmt, setup) 743 | for stmt in stmts 744 | ] 745 | 746 | # Check for exceptions 747 | for timer in timers: 748 | timer.timeit(1) 749 | 750 | # Compare results 751 | def get_result(stmt): 752 | ns = {} 753 | exec(setup + stmt, ns) 754 | return ns.get('result') 755 | 756 | unittest.TestCase('__init__').assertEqual( 757 | *map(get_result, stmts), 758 | '\n=====\nThe two methods yielded different results!' 759 | ) 760 | 761 | # determine number so that 1 <= total time < 3 762 | number = 1 763 | for i in range(22): 764 | number = 3 ** i 765 | if timers[0].timeit(number) >= 1: 766 | break 767 | 768 | print('Number of trials:', number) 769 | print() 770 | 771 | def print_time(idx, el): 772 | print('Method {}: {:.3f}'.format( 773 | idx + 1, el)) 774 | 775 | times = [[] for _ in timers] 776 | for _ in range(repeat): 777 | for i, timer in enumerate(timers): 778 | elapsed = timer.timeit(number) 779 | print_time(i, elapsed) 780 | times[i].append(elapsed) 781 | print() 782 | 783 | print('Best times:') 784 | print('-----------') 785 | for i, elapsed_list in enumerate(times): 786 | print_time(i, min(elapsed_list)) 787 | --------------------------------------------------------------------------------