├── .gitignore ├── HISTORY.md ├── LICENSE ├── MANIFEST.in ├── Makefile ├── README.md ├── aioutils ├── __init__.py ├── bag.py ├── pool.py └── yielder.py ├── setup.py └── tests ├── test_bag.py ├── test_pool.py ├── test_threading.py └── test_yielder.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | # C extensions 6 | *.so 7 | 8 | # Distribution / packaging 9 | .Python 10 | env/ 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | lib/ 17 | lib64/ 18 | parts/ 19 | sdist/ 20 | var/ 21 | *.egg-info/ 22 | .installed.cfg 23 | *.egg 24 | 25 | # PyInstaller 26 | # Usually these files are written by a python script from a template 27 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 28 | *.manifest 29 | *.spec 30 | 31 | # Installer logs 32 | pip-log.txt 33 | pip-delete-this-directory.txt 34 | 35 | # Unit test / coverage reports 36 | htmlcov/ 37 | .tox/ 38 | .coverage 39 | .cache 40 | nosetests.xml 41 | coverage.xml 42 | 43 | # Translations 44 | *.mo 45 | *.pot 46 | 47 | # Django stuff: 48 | *.log 49 | 50 | # Sphinx documentation 51 | docs/_build/ 52 | 53 | # PyBuilder 54 | target/ 55 | 56 | *.swp 57 | .python-version 58 | -------------------------------------------------------------------------------- /HISTORY.md: -------------------------------------------------------------------------------- 1 | ## History 2 | 3 | ### 2015.03.12 4 | 5 | 0.3.10 release 6 | 7 | - do not raise asyncio.InvalidStateError 8 | 9 | ### 2015.03.11 10 | 11 | 0.3.9 release 12 | 13 | - fix a bug will cause yielder stuck 14 | 15 | ### 2015.03.05 16 | 17 | 0.3.8 release 18 | 19 | - when GeneratorExit cancel all pending tasks 20 | - yielder related status cleanups 21 | 22 | ### 2015.03.02 23 | 24 | 0.3.4-0.3.7 release 25 | 26 | - create new event loop if using in a thread 27 | - fix semaphore order so that loop must exist 28 | - add test to test this behaviour 29 | - cleanup _safe_yield_from that are not needed anymore 30 | - use is_running (how can I not using that!!) 31 | 32 | ### 2015.02.28 33 | 34 | 0.3.3 release 35 | 36 | - fixed Pool's potential thread unsafe problem 37 | - fixed Yielder hangs when no items to yield 38 | - add pool size argument to yielder family 39 | 40 | ### 2015.02.27 41 | 42 | 0.3.2 release 43 | 44 | - typo fix 45 | 46 | 0.3.1 release 47 | 48 | - add put method for manually yield items 49 | 50 | 0.3.0 release 51 | 52 | - add Yielder and OrderedYielder to replace Bag and OrderedBag 53 | - fix thread safe problem #2 in Yielder and Group mixed usage 54 | 55 | ### 2015.02.26 56 | 57 | 0.2.1 release 58 | 59 | - fix thread unsafe problem 60 | 61 | ### 2015.02.26 62 | 63 | 0.2.0 release 64 | 65 | - add Bag and OrderedBag 66 | - rename to "aioutils" 67 | 68 | ### 2015.02.25 69 | 70 | 0.1.2 release 71 | 72 | - fix some release problems 73 | 74 | 0.1.1 release 75 | 76 | - basic group and pool implemenation 77 | 78 | 79 | ### 2015.02.23 80 | 81 | ideas, prototypes 82 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2015 Jingchao Hu 2 | 3 | Licensed under the Apache License, Version 2.0 (the "License"); 4 | you may not use this file except in compliance with the License. 5 | You may obtain a copy of the License at 6 | 7 | http://www.apache.org/licenses/LICENSE-2.0 8 | 9 | Unless required by applicable law or agreed to in writing, software 10 | distributed under the License is distributed on an "AS IS" BASIS, 11 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 | See the License for the specific language governing permissions and 13 | limitations under the License. 14 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.md LICENSE HISTORY.md tests/*.py 2 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | develop: 2 | python setup.py develop 3 | publish: 4 | python setup.py sdist upload 5 | test: 6 | PYTHONPATH=. nosetests -v --with-coverage --cover-package=aioutils tests/ 7 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Python3 Asyncio Utils 2 | 3 | ## Introduction 4 | 5 | Python3 Asyncio implements an event loop, but 6 | 7 | - it is quite low level, it misses some advanced primitives 8 | - it can only write sync code wrapped in async code helper (`run_in_executer`), not the other way around 9 | 10 | Specifically, to use asyncio, you must write all your code in an async way, write wrappers for all blocking code, and execute them all in a loop 11 | 12 | That's quite obfuscated for many applications. It is much easier to write critical part in async mode while remain transparent to others sync codes. 13 | 14 | To achieve this, here is package I wrote that provides the following primitives. 15 | 16 | - `Group`: a `gevent.pool.Group` alike object, allows you to spawn coroutines and join them later 17 | - `Pool`: a `gevent.poo.Pool` alike object, allows setting concurrency level 18 | - `Yielder`: a helper to write generator with coroutines 19 | - `OrderedYielder`: a helper to write generator with coroutines, and keep yielding order the same as spawning order 20 | 21 | 22 | ## QuickStart 23 | 24 | ### Group 25 | 26 | Simple `Group` Usage. 27 | 28 | ```py 29 | import random 30 | import asyncio 31 | from aioutils import Pool, Group, Yielder, OrderedYielder, yielding, ordered_yielding 32 | 33 | @asyncio.coroutine 34 | def f(c): 35 | yield from asyncio.sleep(random.random()/10) 36 | return c 37 | 38 | chars = 'abcdefg' 39 | 40 | g = Group() 41 | for c in chars: 42 | g.spawn(f(c)) 43 | g.join() 44 | ``` 45 | 46 | You might find that there is no explicit event loop, it looks just like threading or gevent. 47 | 48 | Under the hood, an event loop starts and runs until all tasks in group joins, 49 | then the event loop stops, allows synced calls from outside. 50 | 51 | The event pool may start later again by other method or in orther thread, it is completely thread safe. 52 | 53 | ### Pool 54 | 55 | Sometimes, you want to limit the maximum level of concurrency, you can use `Pool` instead. 56 | 57 | ```py 58 | p = Pool(2) 59 | for c in chars: 60 | p.spawn(f(c)) 61 | p.join() 62 | ``` 63 | 64 | The only differences between `Pool` and `Group` is that a `Pool` initializes with a integer as the limiting concurrency. 65 | 66 | ### Yielder 67 | 68 | If the return value of the spawned coroutines matters to you, use `Yielder` 69 | 70 | ```py 71 | def gen_func(): 72 | y = Yielder() 73 | for c in chars: 74 | y.spawn(f(c)) 75 | yield from y.yielding() 76 | 77 | print(list(gen_func()) 78 | # outputs an unordered version of ['b', 'd', 'c', 'e', 'f', 'g', 'a'] 79 | ``` 80 | 81 | Note that **`Yielder` only captures results that returns something**, i.e. if you spawn a coroutine that doesn't return anything, or returns `None`, `Yielder` will not yield that value. 82 | 83 | You can also use `y.put` method to explicitly declare what items to be yielded. 84 | 85 | Under the hood, `Yielder` runs an event loop until the first non-None-return corotuine completed, then stops the loop and yield. This process is repeated until all spawned coroutines are done. 86 | 87 | You can also use a context manager `yielding` 88 | 89 | ```py 90 | def gen_func2(): 91 | with yielding() as y: 92 | for c in chars: 93 | y.spawn(f(c)) 94 | yield from y 95 | 96 | print(list(gen_func2()) 97 | # outputs an unordered version of ['b', 'd', 'c', 'e', 'f', 'g', 'a'] 98 | ``` 99 | 100 | The `Yielder` and `yielding` are both thread safe. 101 | 102 | ### Sequential "yield from"s 103 | 104 | When using `yielding`, you'd better avoid using sequential "yield from"s when possible, the problem code is as follows 105 | 106 | ```py 107 | @asyncio.coroutine 108 | def f() 109 | t1 = yield from f1() 110 | t2 = yield from f2() 111 | y.put(t1 + t2) 112 | ``` 113 | 114 | This code alone is ok, but if you use it in a loop, spawning thousands of async tasks of "f", then the event loop need to wait for all f1s complete before it can schedule f2. Thus the yielding processing won't yield anything, it will block until all "f1"s done and some "f2"s done too. 115 | 116 | If there're no dependecies between "f1" and "f2", the following modification works just fine 117 | 118 | ```py 119 | @asyncio.coroutine 120 | def f() 121 | t1 = asyncio.async(f1()) 122 | t2 = asyncio.async(f2()) 123 | yield from asyncio.wait([t1, t2]) 124 | y.put(t1.result() + t2.result()) 125 | ``` 126 | 127 | If f2's argument depends on f1's result, you'd better do "yield from f2()" directly inside f1. 128 | 129 | This coding preference holds true for raw asyncio, too. 130 | 131 | ### OrderedYielder 132 | 133 | If you want the return order to be the same as spawn order, `OrderedYielder` 134 | 135 | ```py 136 | def gen_func(): 137 | y = OrderedYielder() 138 | for c in chars: 139 | y.spawn(f(c)) 140 | yield from y.yielding() 141 | 142 | print(list(gen_func()) 143 | # ['a', 'b', 'c', 'd', 'e', 'f', 'g'] 144 | ``` 145 | 146 | And also there is `ordered_yielding` works just like `yielding` 147 | 148 | ### Examples 149 | 150 | see [test cases](tests) for example usages. 151 | 152 | 153 | ## Testing 154 | 155 | Install nosetests, coverage, then run the following 156 | 157 | ```bash 158 | make test 159 | 160 | # or 161 | 162 | PYTHONPATH=. nosetests tests/ 163 | ``` 164 | 165 | ## More Examples 166 | 167 | The Group is quite useful in complex asynchronous situations. 168 | 169 | Compare these codes below 170 | 171 | ### Simple Loop 172 | 173 | Even in simplest case, Group can make the below code cleaner 174 | 175 | #### 1. using `asyncio` 176 | 177 | ```py 178 | @asyncio.coroutine 179 | def f1(): 180 | yield from asyncio.sleep(1.0) 181 | 182 | tasks = [asyncio.async(f1()) for _ in range(10)] 183 | asyncio.get_event_loop().run_until_complete(asyncio.wait(tasks)) 184 | ``` 185 | 186 | #### 2. using `Group` 187 | 188 | ```py 189 | from aioutils import Group 190 | 191 | g = Group() 192 | for _ in range(10): 193 | g.spawn(f1()) 194 | g.join() 195 | ``` 196 | 197 | ### Nested Loops 198 | 199 | When you have multiple levels of loops, each depends on some blocking data, your situation is worse in raw asyncio 200 | 201 | #### 1. a sync example 202 | 203 | ```py 204 | def f1(): 205 | time.sleep(0.1) 206 | return random.random() 207 | 208 | def f2(v1): 209 | time.sleep(0.1) 210 | return random.random() 211 | 212 | for _ in range(10): 213 | v1 = f1() 214 | for _ in range(10): 215 | f2(v1) 216 | ``` 217 | 218 | Now we need to wait `f1` 10 times, plus `f2` 100 times, it takes 11 seconds to finish. 219 | 220 | But the optimal time for this problem is only 0.2 seconds, that is 221 | 222 | - let all `f1`s runs and return, 0.1 seconds, 223 | - then let all `f2`s run and return, another 0.1 seconds. 224 | 225 | Let's try to write them in aysncio. 226 | 227 | #### 2. using `aysncio` 228 | 229 | 230 | ##### make f1 and f2 coroutine 231 | 232 | ```py 233 | @asyncio.coroutine 234 | def f1(): 235 | yield from asyncio.sleep(0.1) 236 | return random.random() 237 | 238 | @asyncio.coroutine 239 | def f2(v1): 240 | yield from asyncio.sleep(0.1) 241 | return random.random() + v1 242 | ``` 243 | 244 | ##### asyncio, make the second loop coroutine 245 | 246 | ``` 247 | tasks = [] 248 | 249 | @asyncio.coroutine 250 | def inner_task(v1): 251 | for _ in range(10): 252 | yield from f2(v1) 253 | 254 | for _ in range(10): 255 | tasks.append(asyncio.async(inner_task)) 256 | 257 | asyncio.get_event_loop().run_until_complete(asyncio.wait(tasks)) 258 | ``` 259 | 260 | 261 | But this only parallels on the first loop, not the seconds one. 262 | 263 | We need to create two levels of tasks, level1 tasks just runs and generates level2 tasks, then we wait on all level2 tasks. 264 | 265 | ##### asyncio, with 2 task levels 266 | 267 | ```py 268 | level1 = [] 269 | level2 = [] 270 | 271 | @asyncio.coroutine 272 | def inner_task(v1): 273 | for _ in range(10) 274 | level2.append(asyncio.async(f2(v1))) 275 | 276 | for _ in range(10): 277 | level1.append(asyncio.async(inner_task)) 278 | 279 | loop = asyncio.get_event_loop() 280 | loop.run_until_complete(asyncio.wait(level1)) 281 | loop.run_until_complete(asyncio.wait(level2)) 282 | ``` 283 | 284 | We must keep two levels of tasks, otherwise `asyncio.wait` will panic when you try to add new async tasks when level1 tasks got resolved. 285 | 286 | This implementation is not only ugly, and also hard to manage, what if you have 3, 4, or unknown level of tasks? 287 | 288 | #### 3. using `Group` 289 | 290 | ```py 291 | g = Group() 292 | 293 | @asyncio.coroutine 294 | def inner_task(v1): 295 | for _ in range(10) 296 | g.spawn(f2(v1)) 297 | 298 | for _ in range(10): 299 | g.spawn(inner_task(v1)) 300 | 301 | g.join() 302 | ``` 303 | 304 | Which is (I think) easier to understand and use. 305 | 306 | 307 | -------------------------------------------------------------------------------- /aioutils/__init__.py: -------------------------------------------------------------------------------- 1 | from .pool import Pool, Group 2 | from .bag import Bag, OrderedBag 3 | from .yielder import Yielder, OrderedYielder, yielding, ordered_yielding 4 | 5 | __all__ = ['Pool', 'Group', 'Bag', 'OrderedBag', 6 | 'Yielder', 'OrderedYielder', 'yielding', 'ordered_yielding'] 7 | __version__ = '0.3.10' 8 | -------------------------------------------------------------------------------- /aioutils/bag.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ Bag and OrderedBag 4 | 5 | This is old implementation for yielding helper 6 | You should use Yielder and OrderedYielder instead 7 | """ 8 | import time 9 | import queue 10 | import inspect 11 | import asyncio 12 | import functools 13 | import threading 14 | 15 | from .pool import Group 16 | 17 | 18 | class Bag(object): 19 | 20 | """ Create generator function with coroutines 21 | 22 | A Bag is just a Group, a queue, and a background thread running event loop. 23 | Coroutines are spawned in a schedule method, and enqueues result to the 24 | queue, main thread then yield items from queue. 25 | 26 | Usage:: 27 | 28 | def gen_func(): 29 | b = Bag() 30 | @asyncio.coroutine 31 | def coro(arg): 32 | ... 33 | b.put(...) 34 | 35 | def schedule(): 36 | ... 37 | for ...: 38 | b.spawn(coro(arg)) 39 | ... 40 | b.join() 41 | 42 | b.schedule(schedule) 43 | yield from b.yielder() 44 | """ 45 | 46 | def __init__(self, group=None): 47 | if group is not None: 48 | self.g = group 49 | self.loop = self.g.loop 50 | else: 51 | try: 52 | self.loop = asyncio.get_event_loop() 53 | if self.loop.is_running(): 54 | raise NotImplementedError("Cannot use aioutils in " 55 | "asynchroneous environment") 56 | except: 57 | self.loop = asyncio.new_event_loop() 58 | asyncio.set_event_loop(self.loop) 59 | self.g = Group(loop=self.loop) 60 | self.q = queue.Queue() 61 | self.t = None 62 | 63 | def spawn(self, coro): 64 | return self.g.spawn(coro) 65 | 66 | def join(self): 67 | return self.g.join() 68 | 69 | def put(self, item): 70 | self.q.put(item) 71 | 72 | def schedule(self, schedule): 73 | def schedule_wrapper(loop): 74 | asyncio.set_event_loop(loop) 75 | schedule() 76 | 77 | self.t = threading.Thread(target=schedule_wrapper, args=(self.loop,)) 78 | self.t.start() 79 | 80 | def yielder(self): 81 | while (self.t and self.t.is_alive()) or (self.q.qsize() > 0): 82 | try: 83 | yield self.q.get_nowait() 84 | except: 85 | time.sleep(0.1) 86 | 87 | 88 | class OrderedBag(Bag): 89 | 90 | """ A Bag that ensures ordering """ 91 | 92 | def __init__(self, *args): 93 | super(OrderedBag, self).__init__(*args) 94 | self.q = queue.PriorityQueue() 95 | self.order = 0 96 | 97 | def spawn(self, coro): 98 | self.order += 1 99 | 100 | def _coro_wrapper(order): 101 | yield from coro 102 | return self.g.spawn(functools.partial(_coro_wrapper, self.order)()) 103 | 104 | def put(self, item): 105 | order = self._get_coro_order() 106 | self.q.put((order, item)) 107 | 108 | def _get_coro_order(self): 109 | """ Get the coro (and its order) of self """ 110 | stack = inspect.stack() 111 | for frame, module, line, function, context, index in stack: 112 | if function == '_coro_wrapper': 113 | return frame.f_locals['order'] 114 | 115 | def yielder(self): 116 | next_order = 1 117 | while (self.t and self.t.is_alive()) or (self.q.qsize() > 0): 118 | try: 119 | # peek item, the smallest of heap 120 | order, item = self.q.queue[0] 121 | if order != next_order: 122 | time.sleep(0.03) 123 | else: 124 | order, item = self.q.get_nowait() 125 | next_order += 1 126 | yield item 127 | except IndexError: 128 | time.sleep(0.03) 129 | -------------------------------------------------------------------------------- /aioutils/pool.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ Gevent pool/group alike: make asyncio easier to use 4 | 5 | Usage:: 6 | 7 | >>> @asyncio.coroutine 8 | >>> def f(url): 9 | ... r = yield from aiohttp.request('get', url) 10 | ... content = yield from r.read() 11 | ... print('{}: {}'.format(url, content[:80])) 12 | 13 | >>> g = Group() 14 | >>> g.async(f('http://www.baidu.com')) 15 | >>> g.async(f('http://www.sina.com.cn')) 16 | >>> g.join() 17 | 18 | >>> # limit the concurrent coroutines to 3 19 | >>> p = Pool(3) 20 | >>> for _ in range(10): 21 | ... p.async(f('http://www.baidu.com')) 22 | >>> p.join() 23 | """ 24 | import asyncio 25 | 26 | 27 | class Group(object): 28 | 29 | def __init__(self, loop=None): 30 | try: 31 | self.loop = loop or asyncio.get_event_loop() 32 | if self.loop.is_running(): 33 | raise NotImplementedError("Cannot use aioutils in " 34 | "asynchroneous environment") 35 | except: 36 | self.loop = asyncio.new_event_loop() 37 | asyncio.set_event_loop(self.loop) 38 | self._prepare() 39 | 40 | def _prepare(self): 41 | self.counter = 0 42 | self.task_waiter = asyncio.futures.Future(loop=self.loop) 43 | 44 | def spawn(self, coro_or_future): 45 | self.counter += 1 46 | task = asyncio.async(coro_or_future) 47 | task.add_done_callback(self._on_completion) 48 | return task 49 | 50 | async = spawn 51 | 52 | def _on_completion(self, f): 53 | self.counter -= 1 54 | f.remove_done_callback(self._on_completion) 55 | if self.counter <= 0: 56 | if not self.task_waiter.done(): 57 | self.task_waiter.set_result(None) 58 | 59 | def join(self): 60 | def _on_waiter(f): 61 | self.loop.stop() 62 | self._prepare() 63 | 64 | self.task_waiter.add_done_callback(_on_waiter) 65 | 66 | # expect the loops to be stop and start multiple times 67 | while self.counter > 0: 68 | if not self.loop.is_running(): 69 | self.loop.run_forever() 70 | 71 | 72 | class Pool(Group): 73 | 74 | def __init__(self, pool_size, loop=None): 75 | self.sem = asyncio.Semaphore(pool_size, loop=loop) 76 | super(Pool, self).__init__(loop) 77 | 78 | def spawn(self, coro): 79 | assert asyncio.iscoroutine(coro), 'pool only accepts coroutine' 80 | 81 | @asyncio.coroutine 82 | def _limit_coro(): 83 | with (yield from self.sem): 84 | return (yield from coro) 85 | 86 | self.counter += 1 87 | task = asyncio.async(_limit_coro()) 88 | task.add_done_callback(self._on_completion) 89 | return task 90 | 91 | async = spawn 92 | -------------------------------------------------------------------------------- /aioutils/yielder.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | """ Bag and OrderedBag """ 4 | import heapq 5 | import asyncio 6 | import functools 7 | import collections 8 | 9 | 10 | class Yielder(object): 11 | 12 | """ A Bag Rewrite 13 | 14 | - easier to use 15 | - no background threading 16 | 17 | Each time when an item is put, we stop the main loop(!!) and yield 18 | """ 19 | 20 | def __init__(self, pool_size=None): 21 | try: 22 | self.loop = asyncio.get_event_loop() 23 | if self.loop.is_running(): 24 | raise NotImplementedError("Cannot use aioutils in " 25 | "asynchroneous environment") 26 | except: 27 | self.loop = asyncio.new_event_loop() 28 | asyncio.set_event_loop(self.loop) 29 | self.sem = asyncio.Semaphore(pool_size) if pool_size else None 30 | self._prepare() 31 | 32 | def _prepare(self): 33 | self.counter = 0 34 | self.done = collections.deque() 35 | self.getters = collections.deque() 36 | self.exceptions = [] 37 | self.tasks = [] 38 | 39 | def spawn(self, coro): 40 | task = self._async_task(coro) 41 | task.add_done_callback(self._on_completion) 42 | self.counter += 1 43 | self.tasks.append(task) 44 | return task 45 | 46 | def _async_task(self, coro): 47 | if self.sem: 48 | @asyncio.coroutine 49 | def _limit_coro(): 50 | with (yield from self.sem): 51 | return (yield from coro) 52 | task = asyncio.async(_limit_coro()) 53 | else: 54 | task = asyncio.async(coro) 55 | return task 56 | 57 | def _on_completion(self, f): 58 | self.counter -= 1 59 | f.remove_done_callback(self._on_completion) 60 | try: 61 | result = f.result() 62 | except Exception as e: 63 | if not isinstance(e, asyncio.InvalidStateError): 64 | self.exceptions.append(e) 65 | result = None 66 | if result is not None: 67 | self._put(result) 68 | if self.counter <= 0 and self.getters: 69 | getter = self.getters.popleft() 70 | getter.set_result(None) 71 | 72 | def put(self, item): 73 | self._put(item) 74 | 75 | def _put(self, item): 76 | if self.getters: 77 | getter = self.getters.popleft() 78 | getter.set_result(None) 79 | 80 | self.done.append(item) 81 | 82 | def _stop_loop(self, f): 83 | self.loop.stop() 84 | 85 | def _yielding(self): 86 | while self.counter > 0 or self.done: 87 | if self.done: 88 | yield self.done.popleft() 89 | else: 90 | getter = asyncio.Future() 91 | self.getters.append(getter) 92 | getter.add_done_callback(self._stop_loop) 93 | if not self.loop.is_running(): 94 | self.loop.run_forever() 95 | 96 | def yielding(self): 97 | try: 98 | for x in self._yielding(): 99 | if isinstance(x, asyncio.Future): 100 | continue 101 | yield x 102 | except GeneratorExit: 103 | for task in self.tasks: 104 | if not task.done(): 105 | task.cancel() 106 | self.counter -= 1 107 | 108 | if self.exceptions: 109 | raise self.exceptions[0] 110 | self._prepare() 111 | 112 | 113 | class OrderedYielder(Yielder): 114 | def __init__(self, pool_size=None): 115 | super(OrderedYielder, self).__init__(pool_size) 116 | self._prepare() 117 | 118 | def _prepare(self): 119 | super(OrderedYielder, self)._prepare() 120 | self.done = [] 121 | self.order = 0 122 | self.yield_counter = 0 123 | 124 | def spawn(self, coro): 125 | self.order += 1 126 | task = self._async_task(coro) 127 | task.add_done_callback( 128 | functools.partial(self._on_completion, order=self.order)) 129 | self.counter += 1 130 | self.tasks.append(task) 131 | return task 132 | 133 | def _on_completion(self, f, order): 134 | self.counter -= 1 135 | f.remove_done_callback(self._on_completion) 136 | try: 137 | result = f.result() 138 | except Exception as e: 139 | if not isinstance(e, asyncio.InvalidStateError): 140 | self.exceptions.append(e) 141 | result = None 142 | self._put((order, result)) 143 | 144 | def _put(self, item, heappush=heapq.heappush): 145 | order, item = item 146 | if self.yield_counter + 1 == order or self.counter <= 0: 147 | if self.getters: 148 | getter = self.getters.popleft() 149 | getter.set_result(None) 150 | 151 | heappush(self.done, (order, item)) 152 | 153 | def put(self, item): 154 | self.order += 1 155 | self._put((self.order, item)) 156 | 157 | def _yielding(self, heappop=heapq.heappop): 158 | self.yield_counter = 1 159 | while self.counter > 0 or self.done: 160 | if self.done: 161 | order, item = self.done[0] 162 | if self.yield_counter == order: 163 | _, item = heappop(self.done) 164 | if item is not None: 165 | yield item 166 | self.yield_counter += 1 167 | continue 168 | 169 | getter = asyncio.Future() 170 | self.getters.append(getter) 171 | getter.add_done_callback(self._stop_loop) 172 | if not self.loop.is_running(): 173 | self.loop.run_forever() 174 | 175 | 176 | class YieldingContext(object): 177 | def __init__(self, pool_size=None, ordered=False): 178 | if ordered: 179 | self.y = OrderedYielder(pool_size) 180 | else: 181 | self.y = Yielder(pool_size) 182 | self.yielding = None 183 | 184 | def spawn(self, coro): 185 | return self.y.spawn(coro) 186 | 187 | def put(self, item): 188 | return self.y.put(item) 189 | 190 | def __enter__(self): 191 | return iter(self) 192 | 193 | def __iter__(self): 194 | self.yielding = self.y.yielding() 195 | return self 196 | 197 | def __next__(self): 198 | return next(self.yielding) 199 | 200 | def __exit__(self, *args): 201 | pass 202 | 203 | yielding = YieldingContext 204 | ordered_yielding = functools.partial(YieldingContext, ordered=True) 205 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | from setuptools import setup 4 | 5 | from aioutils import __version__ 6 | 7 | with open('HISTORY.md') as f: 8 | history = f.read() 9 | 10 | setup( 11 | name='aioutils', 12 | version=__version__, 13 | description='Python3 Asyncio Utils', 14 | long_description= history, 15 | author='Jingchao Hu', 16 | author_email='jingchaohu@gmail.com', 17 | url='https://github.com/observerss/aioutils', 18 | packages=['aioutils'], 19 | license='Apache 2.0', 20 | zip_safe=True, 21 | classifiers=[ 22 | 'Development Status :: 5 - Production/Stable', 23 | 'Intended Audience :: Developers', 24 | 'Natural Language :: English', 25 | 'License :: OSI Approved :: Apache Software License', 26 | 'Programming Language :: Python :: 3', 27 | 'Programming Language :: Python :: 3.4' 28 | ], 29 | ) 30 | -------------------------------------------------------------------------------- /tests/test_bag.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import random 4 | import inspect 5 | import asyncio 6 | 7 | from aioutils import Bag, OrderedBag, Group 8 | 9 | def test_bag(): 10 | chars = 'abcdefg' 11 | def g(): 12 | b = Bag() 13 | @asyncio.coroutine 14 | def f(c): 15 | yield from asyncio.sleep(random.random()/10) 16 | b.put(c) 17 | def schedule(): 18 | for c in chars: 19 | b.spawn(f(c)) 20 | b.join() 21 | b.schedule(schedule) 22 | yield from b.yielder() 23 | 24 | chars2 = g() 25 | assert inspect.isgenerator(chars2) 26 | chars2 = list(chars2) 27 | assert set(chars) == set(chars2) 28 | 29 | 30 | def test_orderedbag(): 31 | chars = 'abcdefg' 32 | def g(): 33 | b = OrderedBag(Group()) 34 | @asyncio.coroutine 35 | def f(c): 36 | yield from asyncio.sleep(random.random()*0.1) 37 | b.put(c) 38 | def schedule(): 39 | for c in chars: 40 | b.spawn(f(c)) 41 | b.join() 42 | b.schedule(schedule) 43 | yield from b.yielder() 44 | 45 | chars2 = g() 46 | assert inspect.isgenerator(chars2) 47 | chars2 = list(chars2) 48 | for c1, c2 in zip(chars, chars2): 49 | assert c1 == c2 50 | 51 | 52 | if __name__ == '__main__': 53 | test_bag() 54 | test_orderedbag() 55 | -------------------------------------------------------------------------------- /tests/test_pool.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import time 4 | import asyncio 5 | 6 | from aioutils import Group, Pool 7 | 8 | def test_group(): 9 | timespan = 0.1 10 | @asyncio.coroutine 11 | def f(i): 12 | t0 = time.time() 13 | yield from asyncio.sleep(timespan) 14 | t = time.time() - t0 15 | print('finish {}, seconds={:4.2f}'.format(i, t)) 16 | 17 | print('testing group') 18 | t0 = time.time() 19 | g = Group() 20 | for i in range(9): 21 | g.spawn(f(i)) 22 | g.join() 23 | print('total time: {:4.2f}'.format(time.time() - t0)) 24 | assert timespan < time.time() - t0 < timespan * 1.1 25 | 26 | 27 | def test_pool(): 28 | timespan = 0.1 29 | 30 | @asyncio.coroutine 31 | def f(i): 32 | t0 = time.time() 33 | yield from asyncio.sleep(timespan) 34 | t = time.time() - t0 35 | print('finish {}, seconds={:4.2f}'.format(i, t)) 36 | 37 | print('testing pool') 38 | t0 = time.time() 39 | p = Pool(3) 40 | for i in range(9): 41 | p.spawn(f(i)) 42 | p.join() 43 | print('total time: {:4.2f}'.format(time.time() - t0)) 44 | assert timespan * 3 < time.time() - t0 < timespan * 3 * 1.1 45 | 46 | if __name__ == '__main__': 47 | test_group() 48 | test_pool() 49 | -------------------------------------------------------------------------------- /tests/test_threading.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import time 4 | import random 5 | import asyncio 6 | import threading 7 | from aioutils import Group, Yielder, yielding 8 | 9 | @asyncio.coroutine 10 | def f(c): 11 | yield from asyncio.sleep(random.random()*0.02) 12 | return c 13 | 14 | 15 | def test_group_threading(): 16 | """ Ensure that Pool and Group are thread-safe """ 17 | stopall = False 18 | def t(): 19 | while not stopall: 20 | g = Group() 21 | for i in range(5): 22 | g.spawn(f(i)) 23 | 24 | g.join() 25 | 26 | time.sleep(random.random()*0.02) 27 | 28 | tasks = [threading.Thread(target=t) for _ in range(5)] 29 | for task in tasks: task.daemon = True 30 | for task in tasks: task.start() 31 | time.sleep(0.2) 32 | stopall = True 33 | for task in tasks: task.join() 34 | assert asyncio.Task.all_tasks() == set(), asyncio.Task.all_tasks() 35 | 36 | 37 | def test_yielder_threading(): 38 | """ Ensure Yielder are thread safe """ 39 | stopall = False 40 | chars = 'abcdefg' 41 | 42 | def gen_func(): 43 | y = Yielder() 44 | 45 | for c in chars: 46 | y.spawn(f(c)) 47 | 48 | yield from y.yielding() 49 | 50 | def t(): 51 | while not stopall: 52 | chars2 = list(gen_func()) 53 | assert set(chars2) == set(chars) 54 | 55 | time.sleep(random.random()*0.02) 56 | 57 | tasks = [threading.Thread(target=t) for _ in range(5)] 58 | for task in tasks: task.daemon = True 59 | for task in tasks: task.start() 60 | time.sleep(0.2) 61 | stopall = True 62 | for task in tasks: task.join() 63 | assert asyncio.Task.all_tasks() == set(), asyncio.Task.all_tasks() 64 | 65 | 66 | def test_mixed(): 67 | """ Ensure mixed usage are thread safe """ 68 | chars = 'abcdefg' 69 | stopall = False 70 | 71 | def f1(): 72 | y = Yielder() 73 | for c in chars: 74 | y.spawn(f(c)) 75 | 76 | return list(y.yielding()) 77 | 78 | def f2(): 79 | g = Group() 80 | for c in chars: 81 | g.spawn(f(c)) 82 | 83 | g.join() 84 | 85 | def t(): 86 | while not stopall: 87 | f = random.choice([f1, f2]) 88 | r = f() 89 | if f == f1: 90 | assert set(r) == set(chars) 91 | time.sleep(random.random()*0.02) 92 | 93 | tasks = [threading.Thread(target=t) for _ in range(5)] 94 | for task in tasks: task.daemon = True 95 | for task in tasks: task.start() 96 | time.sleep(0.2) 97 | stopall = True 98 | for task in tasks: task.join() 99 | assert asyncio.Task.all_tasks() == set(), asyncio.Task.all_tasks() 100 | 101 | 102 | def test_yielding_size_in_threading(): 103 | chars = 'abcdefgh' 104 | def f1(): 105 | with yielding(2) as y: 106 | for c in chars: 107 | y.spawn(f(c)) 108 | 109 | yield from y 110 | 111 | l = [] 112 | def f2(): 113 | for x in f1(): 114 | l.append(x) 115 | 116 | t = threading.Thread(target=f2) 117 | t.start() 118 | t.join() 119 | assert set(l) == set(chars) 120 | 121 | 122 | if __name__ == '__main__': 123 | test_group_threading() 124 | test_yielder_threading() 125 | test_mixed() 126 | test_yielding_size_in_threading() 127 | -------------------------------------------------------------------------------- /tests/test_yielder.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | # -*- coding: utf-8 -*- 3 | import time 4 | import random 5 | import inspect 6 | import asyncio 7 | from aioutils import Yielder, OrderedYielder, yielding, ordered_yielding 8 | 9 | from nose.tools import raises 10 | 11 | @asyncio.coroutine 12 | def f(c): 13 | yield from asyncio.sleep(random.random()*.1) 14 | return c 15 | 16 | 17 | def test_yielder(): 18 | def gen_func(): 19 | y = Yielder() 20 | @asyncio.coroutine 21 | def f1(i): 22 | yield from asyncio.sleep(random.random()*0.1) 23 | for j in [random.randint(0, 10) for _ in range(i)]: 24 | y.spawn(f2(j)) 25 | 26 | @asyncio.coroutine 27 | def f2(j): 28 | yield from asyncio.sleep(random.random()*0.1) 29 | return j + random.random() 30 | 31 | for i in range(5): 32 | y.spawn(f1(i)) 33 | 34 | y.put(3) 35 | 36 | for i in range(5, 10): 37 | y.spawn(f1(i)) 38 | 39 | yield from y.yielding() 40 | 41 | g = gen_func() 42 | assert inspect.isgenerator(g) 43 | t0 = time.time() 44 | gs = list(g) 45 | assert len(gs) == 46 46 | assert min(gs) > 0 and max(gs) < 11 47 | assert time.time() - t0 < 0.2 * 1.1 48 | 49 | 50 | def test_ordered_yielder(): 51 | chars = 'abcdefg' 52 | 53 | def gen_func(): 54 | y = OrderedYielder() 55 | for c in chars[:3]: 56 | y.spawn(f(c)) 57 | 58 | y.put('z') 59 | 60 | for c in chars[3:]: 61 | y.spawn(f(c)) 62 | 63 | yield from y.yielding() 64 | 65 | g = gen_func() 66 | assert inspect.isgenerator(g) 67 | t0 = time.time() 68 | gs = list(g) 69 | assert ''.join(gs) == chars[:3] + 'z' + chars[3:] 70 | assert time.time() - t0 < 0.1 * 1.1 71 | 72 | 73 | def test_yielding(): 74 | chars = 'abcdefg' 75 | def gen_func(): 76 | with yielding() as y: 77 | for c in chars: 78 | y.spawn(f(c)) 79 | yield from y 80 | 81 | def gen_func2(): 82 | with ordered_yielding() as y: 83 | for c in chars: 84 | y.spawn(f(c)) 85 | yield from y 86 | 87 | g = gen_func() 88 | assert inspect.isgenerator(g) 89 | t0 = time.time() 90 | gs = list(g) 91 | assert set(gs) == set(chars) 92 | assert time.time() - t0 < 0.1 * 1.1 93 | 94 | assert ''.join(list(gen_func2())) == chars 95 | 96 | 97 | def test_yielder_with_pool_size(): 98 | chars = 'abcdefgh' 99 | def gen_func(): 100 | with yielding(2) as y: 101 | for c in chars: 102 | y.spawn(f(c)) 103 | yield from y 104 | 105 | 106 | # I don't know how to assert correctness, >.< 107 | assert set(gen_func()) == set(chars) 108 | 109 | 110 | def test_empty_yielder(): 111 | def gen_func(): 112 | with yielding() as y: 113 | y.spawn(asyncio.sleep(0.01)) 114 | yield from y 115 | 116 | assert list(gen_func()) == [] 117 | 118 | 119 | def test_two_level_ordered_yielding(): 120 | chars = 'abcdefg' 121 | def gen_func(): 122 | with ordered_yielding() as y: 123 | @asyncio.coroutine 124 | def g(i): 125 | for c in chars[:i]: 126 | y.spawn(f(c)) 127 | 128 | for i in range(3, 5): 129 | y.spawn(g(i)) 130 | 131 | yield from y 132 | 133 | assert ''.join(list(gen_func())) == 'abcabcd' 134 | 135 | 136 | def test_break_from_yielding(): 137 | @asyncio.coroutine 138 | def g(c): 139 | yield from asyncio.sleep(random.random()*.1) 140 | return c 141 | y = Yielder() 142 | 143 | def gen_func(): 144 | for c in 'abcdefg': 145 | y.spawn(g(c)) 146 | yield from y.yielding() 147 | 148 | for x in gen_func(): 149 | break 150 | 151 | assert y.counter == 0 152 | 153 | 154 | @raises(ValueError) 155 | def test_raise_from_yielding(): 156 | @asyncio.coroutine 157 | def g(c): 158 | yield from asyncio.sleep(random.random()*.1) 159 | if random.random() < 0.2: 160 | raise ValueError 161 | return c 162 | 163 | def gen_func(): 164 | with yielding(3) as y: 165 | for c in 'abcdefghijklmn': 166 | y.spawn(g(c)) 167 | yield from y 168 | 169 | # test that raise will not cause problem 170 | for x in gen_func(): 171 | pass 172 | 173 | 174 | @raises(ValueError) 175 | def test_raise_from_nested_yielding(): 176 | yielder = Yielder() 177 | 178 | @asyncio.coroutine 179 | def worker(err_count=0): 180 | yield from asyncio.sleep(random.random()*.02) 181 | if err_count > 5: 182 | raise ValueError 183 | else: 184 | return (yield from worker(err_count+1)) 185 | 186 | for i in range(10): 187 | yielder.spawn(worker()) 188 | 189 | for i, ch in enumerate(yielder.yielding()): 190 | pass 191 | 192 | 193 | if __name__ == '__main__': 194 | test_yielder() 195 | test_ordered_yielder() 196 | test_yielding() 197 | test_yielder_with_pool_size() 198 | test_empty_yielder() 199 | test_two_level_ordered_yielding() 200 | test_break_from_yielding() 201 | test_raise_from_yielding() 202 | test_raise_from_nested_yielding() 203 | --------------------------------------------------------------------------------