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 

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: 123 ↛ 122line 123 didn't jump to line 122, because the condition on line 123 was never false

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() 

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 

354class ThrottleState(object): 

355 __slots__ = ('running', 'pending', 'args') 

356 

357 running: bool 

358 pending: bool 

359 args: Tuple[tuple, dict] 

360 

361 def __init__(self): 

362 self.running = self.pending = False 

363 

364 

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

366 def decorate(f): 

367 deadline = -1.0 

368 gr = None 

369 

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

371 nonlocal deadline, gr 

372 t = deadline - time() 

373 while t > 0: 

374 gevent.sleep(t) 

375 t = deadline - time() 

376 gr = None 

377 deadline = time() + seconds 

378 f(*a, **k) 

379 

380 @wraps(f) 

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

382 nonlocal deadline, gr 

383 t = deadline - time() 

384 if t < 0: 

385 f(*a, **k) 

386 deadline = time() + seconds 

387 else: 

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

389 

390 wrapper.__name__ == f.__name__ 

391 return wrapper 

392 

393 return decorate 

394 

395 

396class InstanceHookMeta(type): 

397 # ABCMeta would use __subclasshook__ for instance check. Loses information. 

398 

399 def __instancecheck__(cls, inst): 

400 return cls.instancecheck(inst) 

401 

402 def __subclasscheck__(cls, C): 

403 return cls.subclasscheck(C) 

404 

405 def instancecheck(cls, inst): 

406 return cls.subclasscheck(type(inst)) 

407 

408 

409class ArgValidationError(Exception): 

410 pass 

411 

412 

413class ArgTypeError(ArgValidationError): 

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

415 

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

417 self.position = position 

418 self.expected = expected 

419 self.actual = actual 

420 

421 def __unicode__(self): 

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

423 self.position, 

424 self.expected.__name__, 

425 self.actual.__name__, 

426 ) 

427 

428 def __str__(self): 

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

430 

431 

432class ArgCountError(ArgValidationError): 

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

434 

435 def __init__(self, expected, actual): 

436 self.expected = expected 

437 self.actual = actual 

438 

439 def __unicode__(self): 

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

441 self.expected, 

442 self.actual, 

443 ) 

444 

445 def __str__(self): 

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

447 

448 

449def validate_args(*typelist): 

450 def decorate(f): 

451 @wraps(f) 

452 def wrapper(*args): 

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

454 if e != a: 

455 raise ArgCountError(e, a) 

456 

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

458 if not isinstance(v, e): 

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

460 

461 return f(*args) 

462 

463 wrapper.__name__ = f.__name__ 

464 return wrapper 

465 

466 return decorate 

467 

468 

469class BusinessException(Exception): 

470 snake_case: str 

471 

472 

473class BusinessExceptionGenerator(object): 

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

475 snake_case = '_'.join([ 

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

477 ]) 

478 

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

480 setattr(self, k, cls) 

481 return cls 

482 

483 

484exceptions = BusinessExceptionGenerator() 

485 

486 

487# Helper to make mypy happy 

488class MockMeta(type): 

489 def mro(cls): 

490 bases = cls.__bases__ 

491 assert len(bases) == 1 

492 return [cls, object]