1# -*- coding: utf-8 -*- 

2from __future__ import annotations 

3 

4# -- stdlib -- 

5from collections import deque 

6from contextlib import contextmanager 

7from functools import wraps 

8from time import time 

9from typing import Any, Callable, Dict, Iterable, List, Set, Tuple, Type, TypeVar 

10from weakref import WeakSet 

11import functools 

12import logging 

13import re 

14 

15# -- third party -- 

16from gevent.lock import Semaphore 

17from gevent.queue import Queue 

18import gevent 

19 

20# -- own -- 

21 

22# -- code -- 

23log = logging.getLogger('util.misc') 

24dbgvals: Dict[str, object] = {} 

25 

26 

27class ObjectDict(dict): 

28 __slots__ = () 

29 

30 def __getattr__(self, k): 

31 try: 

32 return self[k] 

33 except KeyError: 

34 raise AttributeError 

35 

36 def __setattr__(self, k, v): 

37 self[k] = v 

38 

39 @classmethod 

40 def parse(cls, data): 

41 if isinstance(data, dict): 

42 return cls({k: cls.parse(v) for k, v in data.items()}) 

43 elif isinstance(data, (list, tuple, set, frozenset)): 

44 return type(data)([cls.parse(v) for v in data]) 

45 

46 return data 

47 

48 

49T = TypeVar('T') 

50 

51 

52class BatchList(List[T]): 

53 __slots__ = () 

54 

55 def __getattribute__(self, name): 

56 try: 

57 list_attr = list.__getattribute__(self, name) 

58 return list_attr 

59 except AttributeError: 

60 pass 

61 

62 return BatchList( 

63 getattr(i, name) for i in self 

64 ) 

65 

66 def __call__(self, *a, **k): 

67 return BatchList( 

68 f(*a, **k) for f in self # type: ignore 

69 ) 

70 

71 ''' 

72 @typing.overload 

73 def __getitem__(self, i: int) -> T: ... 

74 

75 @typing.overload 

76 def __getitem__(self, s: slice) -> 'BatchList[T]': ... 

77 

78 def __getitem__(self, a): 

79 if isinstance(a, slice): 

80 return BatchList(list.__getitem__(self, a)) 

81 else: 

82 return list.__getitem__(self, a) 

83 ''' 

84 

85 def exclude(self, elem: T) -> 'BatchList[T]': 

86 nl = BatchList(self) 

87 try: 

88 nl.remove(elem) 

89 except ValueError: 

90 pass 

91 

92 return nl 

93 

94 def rotate_to(self, elem: T) -> 'BatchList[T]': 

95 i = self.index(elem) 

96 n = len(self) 

97 return self.__class__((self*2)[i:i+n]) 

98 

99 def replace(self, old: T, new: T) -> bool: 

100 try: 

101 self[self.index(old)] = new 

102 return True 

103 except ValueError: 

104 return False 

105 

106 def find_replace(self, pred: Callable[[T], bool], new: T) -> bool: 

107 for i, v in enumerate(self): 

108 if pred(v): 

109 self[i] = new 

110 return True 

111 else: 

112 return False 

113 

114 def sibling(self, me: T, offset=1) -> T: 

115 i = self.index(me) 

116 n = len(self) 

117 return self[(i + offset) % n] 

118 

119 

120def remove_dups(s): 

121 seen: Set[Any] = set() 

122 for i in s: 

123 if i not in seen: 

124 yield i 

125 seen.add(i) 

126 

127 

128def classmix(*_classes) -> type: 

129 classes: Any = [] 

130 for c in _classes: 

131 if hasattr(c, '_is_mixedclass'): 

132 classes.extend(c.__bases__) 

133 else: 

134 classes.append(c) 

135 

136 classes = tuple(remove_dups(classes)) 

137 cached = cls_cache.get(classes, None) 

138 if cached: return cached 

139 

140 clsname = ', '.join(cls.__name__ for cls in classes) 

141 new_cls = type('Mixed(%s)' % clsname, classes, {'_is_mixedclass': True}) 

142 cls_cache[classes] = new_cls 

143 return new_cls 

144 

145 

146cls_cache: Dict[tuple, type] = {} 

147 

148 

149def hook(module): 

150 def inner(hooker): 

151 funcname = hooker.__name__ 

152 hookee = getattr(module, funcname) 

153 

154 @wraps(hookee) 

155 def real_hooker(*args, **kwargs): 

156 return hooker(hookee, *args, **kwargs) 

157 setattr(module, funcname, real_hooker) 

158 return real_hooker 

159 return inner 

160 

161 

162def extendclass(clsname, bases, _dict): 

163 for cls in bases: 

164 for key, value in _dict.items(): 

165 if key == '__module__': 

166 continue 

167 setattr(cls, key, value) 

168 

169 

170def partition(pred: Callable[[Any], bool], lst: Iterable[Any]) -> Tuple[List[Any], List[Any]]: 

171 f: List[Any] 

172 t: List[Any] 

173 f, t = [], [] 

174 for i in lst: 

175 (f, t)[pred(i)].append(i) 

176 

177 return t, f 

178 

179 

180def track(f): 

181 @functools.wraps(f) 

182 def _wrapper(*a, **k): 

183 print('%s: %s %s' % (f.__name__, a, k)) 

184 return f(*a, **k) 

185 return _wrapper 

186 

187 

188def flatten(l): 

189 rst = [] 

190 

191 def _flatten(sl): 

192 for i in sl: 

193 if isinstance(i, (list, tuple, deque)): 

194 _flatten(i) 

195 else: 

196 rst.append(i) 

197 

198 _flatten(l) 

199 return rst 

200 

201 

202def group_by(l, keyfunc): 

203 if not l: return [] 

204 

205 grouped = [] 

206 group = [] 

207 

208 lastkey = keyfunc(l[0]) 

209 for i in l: 

210 k = keyfunc(i) 

211 if k == lastkey: 

212 group.append(i) 

213 else: 

214 grouped.append(group) 

215 group = [i] 

216 lastkey = k 

217 

218 if group: 

219 grouped.append(group) 

220 

221 return grouped 

222 

223 

224def instantiate(cls): 

225 return cls() 

226 

227 

228def surpress_and_restart(f): 

229 def wrapper(*a, **k): 

230 while True: 

231 try: 

232 return f(*a, **k) 

233 except Exception as e: 

234 import logging 

235 log = logging.getLogger('misc') 

236 log.exception(e) 

237 

238 return wrapper 

239 

240 

241def swallow(f): 

242 def wrapper(*a, **k): 

243 try: 

244 return f(*a, **k) 

245 except Exception: 

246 pass 

247 

248 return wrapper 

249 

250 

251def log_failure(logger): 

252 def decorate(f): 

253 def wrapper(*a, **k): 

254 try: 

255 return f(*a, **k) 

256 except Exception as e: 

257 logger.exception(e) 

258 raise 

259 

260 return wrapper 

261 

262 return decorate 

263 

264 

265class ObservableEvent(object): 

266 listeners: Set[Callable] 

267 

268 def __init__(self, weakref=False): 

269 self.listeners = WeakSet() if weakref else set() # type: ignore 

270 

271 def __iadd__(self, ob): 

272 self.listeners.add(ob) 

273 return self 

274 

275 def __isub__(self, ob): 

276 self.listeners.discard(ob) 

277 return self 

278 

279 def notify(self, *a, **k): 

280 for ob in list(self.listeners): 

281 ob(*a, **k) 

282 

283 

284class GenericPool(object): 

285 def __init__(self, factory, size, container_class=Queue): 

286 self.factory = factory 

287 self.size = size 

288 self.container = container_class(size) 

289 self.inited = False 

290 

291 def __call__(self): 

292 @contextmanager 

293 def manager(): 

294 container = self.container 

295 

296 if not self.inited: 

297 for i in range(self.size): 

298 container.put(self.factory()) 

299 

300 self.inited = True 

301 

302 try: 

303 item = container.get() 

304 yield item 

305 except Exception: 

306 item = self.factory() 

307 raise 

308 finally: 

309 try: 

310 container.put_nowait(item) 

311 except Exception: 

312 pass 

313 

314 return manager() 

315 

316 

317def debounce(seconds): 

318 def decorate(f): 

319 lock = Semaphore(1) 

320 last = None 

321 

322 def bouncer(fire, *a, **k): 

323 nonlocal last 

324 gevent.sleep(seconds) 

325 last = None 

326 fire and f(*a, **k) 

327 

328 @wraps(f) 

329 def wrapper(*a, **k): 

330 nonlocal last 

331 rst = lock.acquire(blocking=False) 

332 if not rst: 

333 return 

334 

335 try: 

336 run = False 

337 if last is None: 

338 last = gevent.spawn(bouncer, False) 

339 run = True 

340 else: 

341 last.kill() 

342 last = gevent.spawn(bouncer, True, *a, **k) 

343 finally: 

344 lock.release() 

345 

346 run and f(*a, **k) 

347 

348 wrapper.__name__ == f.__name__ 

349 return wrapper 

350 

351 return decorate 

352 

353 

354DO_NOT_THROTTLE = False 

355 

356 

357def throttle(seconds: float) -> Callable[[T], T]: 

358 def decorate(f): 

359 deadline = -1.0 

360 gr = None 

361 

362 def bouncer(*a, **k): 

363 nonlocal deadline, gr 

364 t = deadline - time() 

365 while t > 0: 

366 gevent.sleep(t) 

367 t = deadline - time() 

368 gr = None 

369 deadline = time() + seconds 

370 f(*a, **k) 

371 

372 @wraps(f) 

373 def wrapper(*a, **k): 

374 if DO_NOT_THROTTLE: 374 ↛ 378line 374 didn't jump to line 378, because the condition on line 374 was never false

375 return f(*a, **k) 

376 

377 nonlocal deadline, gr 

378 t = deadline - time() 

379 if t < 0: 

380 f(*a, **k) 

381 deadline = time() + seconds 

382 else: 

383 gr = gr or gevent.spawn(bouncer, *a, **k) 

384 

385 wrapper.__name__ == f.__name__ 

386 return wrapper 

387 

388 return decorate 

389 

390 

391class ArgValidationError(Exception): 

392 pass 

393 

394 

395class ArgTypeError(ArgValidationError): 

396 __slots__ = ('position', 'expected', 'actual') 

397 

398 def __init__(self, position, expected, actual): 

399 self.position = position 

400 self.expected = expected 

401 self.actual = actual 

402 

403 def __unicode__(self): 

404 return 'Arg %s should be "%s" type, "%s" found' % ( 

405 self.position, 

406 self.expected.__name__, 

407 self.actual.__name__, 

408 ) 

409 

410 def __str__(self): 

411 return self.__unicode__().encode('utf-8') 

412 

413 

414class ArgCountError(ArgValidationError): 

415 __slots__ = ('expected', 'actual') 

416 

417 def __init__(self, expected, actual): 

418 self.expected = expected 

419 self.actual = actual 

420 

421 def __unicode__(self): 

422 return 'Expecting %s args, %s found' % ( 

423 self.expected, 

424 self.actual, 

425 ) 

426 

427 def __str__(self): 

428 return self.__unicode__().encode('utf-8') 

429 

430 

431def validate_args(*typelist): 

432 def decorate(f): 

433 @wraps(f) 

434 def wrapper(*args): 

435 e, a = len(typelist), len(args) 

436 if e != a: 

437 raise ArgCountError(e, a) 

438 

439 for i, e, v in zip(range(1000), typelist, args): 

440 if not isinstance(v, e): 

441 raise ArgValidationError(i, e, v.__class__) 

442 

443 return f(*args) 

444 

445 wrapper.__name__ = f.__name__ 

446 return wrapper 

447 

448 return decorate 

449 

450 

451class BusinessException(Exception): 

452 name: str 

453 snake_case: str 

454 

455 

456class BusinessExceptionGenerator(object): 

457 def __getattr__(self, k: str) -> Type[BusinessException]: 

458 snake_case = '_'.join([ 

459 i.lower() for i in re.findall(r'[A-Z][a-z]+|[A-Z]+(?![a-z])', k) 

460 ]) 

461 

462 cls = type(k, (BusinessException,), {'name': k, 'snake_case': snake_case}) 

463 setattr(self, k, cls) 

464 return cls 

465 

466 

467exceptions = BusinessExceptionGenerator() 

468 

469 

470# Helper to make mypy happy 

471class MockMeta(type): 

472 def mro(cls): 

473 bases = cls.__bases__ 

474 assert len(bases) == 1 

475 return [cls, object] 

476 

477 

478class LoopBreaker(object): 

479 def __init__(self, resets_to=True): 

480 self._should_continue = True 

481 self._reset_to = resets_to 

482 

483 def __iter__(self): 

484 return self 

485 

486 def __next__(self): 

487 if self._should_continue: 

488 self._should_continue = self._reset_to 

489 return self 

490 else: 

491 raise StopIteration 

492 

493 def stop(self): 

494 self._should_continue = False 

495 

496 def cont(self): 

497 self._should_continue = True