refactoring some code
[rainbowstream.git] / rainbowstream / draw.py
CommitLineData
e9f5200b
VNM
1import random
2import itertools
7500d90b
VNM
3import requests
4import datetime
5import time
59262e95 6import re
7500d90b 7
2da50cc4 8from twitter.util import printNicely
e9f5200b
VNM
9from functools import wraps
10from pyfiglet import figlet_format
7500d90b 11from dateutil import parser
2da50cc4 12from .c_image import *
7500d90b
VNM
13from .colors import *
14from .config import *
15from .db import *
c3bab4ef 16from .py3patch import *
17
7500d90b
VNM
18
19db = RainbowDB()
1fdd6a5c
VNM
20g = {}
21
422dd385 22
e9f5200b
VNM
23def init_cycle():
24 """
25 Init the cycle
26 """
27 colors_shuffle = [globals()[i.encode('utf8')]
fa6e062d
O
28 if not str(i).isdigit()
29 else term_color(int(i))
422dd385 30 for i in c['CYCLE_COLOR']]
5f22104f 31 return itertools.cycle(colors_shuffle)
1fdd6a5c 32
e9f5200b 33
59262e95 34def start_cycle():
2359c276
VNM
35 """
36 Notify from rainbow
37 """
5f22104f 38 g['cyc'] = init_cycle()
e43ebfa6 39 g['cache'] = {}
2359c276
VNM
40
41
e9f5200b
VNM
42def order_rainbow(s):
43 """
44 Print a string with ordered color with each character
45 """
5f22104f 46 colors_shuffle = [globals()[i.encode('utf8')]
fa6e062d
O
47 if not str(i).isdigit()
48 else term_color(int(i))
422dd385 49 for i in c['CYCLE_COLOR']]
5f22104f 50 colored = [colors_shuffle[i % 7](s[i]) for i in xrange(len(s))]
c3bab4ef 51 return ''.join(colored)
e9f5200b
VNM
52
53
54def random_rainbow(s):
55 """
56 Print a string with random color with each character
57 """
5f22104f 58 colors_shuffle = [globals()[i.encode('utf8')]
fa6e062d
O
59 if not str(i).isdigit()
60 else term_color(int(i))
422dd385 61 for i in c['CYCLE_COLOR']]
5f22104f 62 colored = [random.choice(colors_shuffle)(i) for i in s]
c3bab4ef 63 return ''.join(colored)
e9f5200b
VNM
64
65
66def Memoize(func):
67 """
68 Memoize decorator
69 """
e9f5200b
VNM
70 @wraps(func)
71 def wrapper(*args):
e43ebfa6 72 if args not in g['cache']:
73 g['cache'][args] = func(*args)
74 return g['cache'][args]
e9f5200b
VNM
75 return wrapper
76
77
78@Memoize
79def cycle_color(s):
80 """
81 Cycle the colors_shuffle
82 """
1fdd6a5c 83 return next(g['cyc'])(s)
e9f5200b
VNM
84
85
86def ascii_art(text):
87 """
88 Draw the Ascii Art
89 """
90 fi = figlet_format(text, font='doom')
91 print('\n'.join(
1fdd6a5c 92 [next(g['cyc'])(i) for i in fi.split('\n')]
e9f5200b
VNM
93 ))
94
95
92685d21 96def check_config():
ceec8593 97 """
98 Check if config is changed
99 """
100 changed = False
101 data = get_all_config()
102 for key in c:
103 if key in data:
104 if data[key] != c[key]:
105 changed = True
106 if changed:
107 reload_config()
108
109
110def validate_theme(theme):
111 """
112 Validate a theme exists or not
113 """
114 # Theme changed check
115 files = os.listdir(os.path.dirname(__file__) + '/colorset')
116 themes = [f.split('.')[0] for f in files if f.split('.')[-1] == 'json']
117 return theme in themes
92685d21 118
119
ceec8593 120def reload_theme(current_config):
4cf86720
VNM
121 """
122 Check current theme and update if necessary
123 """
124 exists = db.theme_query()
125 themes = [t.theme_name for t in exists]
ceec8593 126 if current_config != themes[0]:
1f2f6159 127 config = os.path.dirname(
ceec8593 128 __file__) + '/colorset/' + current_config + '.json'
4cf86720
VNM
129 # Load new config
130 data = load_config(config)
a5301bc0
VNM
131 if data:
132 for d in data:
133 c[d] = data[d]
ceec8593 134 # Restart color cycle and update db/config
135 start_cycle()
136 db.theme_update(current_config)
137 set_config('THEME', current_config)
7500d90b 138
fe08f905
VNM
139
140def color_func(func_name):
141 """
142 Call color function base on name
143 """
fa6e062d
O
144 if str(func_name).isdigit():
145 return term_color(int(func_name))
c3bab4ef 146 return globals()[func_name]
fe08f905
VNM
147
148
fe9bb33b 149def draw(t, keyword=None, check_semaphore=False, fil=[], ig=[]):
7500d90b
VNM
150 """
151 Draw the rainbow
152 """
153
c37c04a9
O
154 # Check the semaphore lock (stream process only)
155 if check_semaphore:
156 if db.semaphore_query_pause():
157 return
158 while db.semaphore_query_lock():
159 time.sleep(0.5)
160
161 # Check config and theme
92685d21 162 check_config()
ceec8593 163 reload_theme(c['THEME'])
c37c04a9 164
7500d90b
VNM
165 # Retrieve tweet
166 tid = t['id']
606def7e 167 text = t['text']
7500d90b
VNM
168 screen_name = t['user']['screen_name']
169 name = t['user']['name']
170 created_at = t['created_at']
171 favorited = t['favorited']
172 date = parser.parse(created_at)
173 date = date - datetime.timedelta(seconds=time.timezone)
8c7ae594
O
174 clock_format = '%Y/%m/%d %H:%M:%S'
175 try:
176 clock_format = c['FORMAT']['TWEET']['CLOCK_FORMAT']
177 except:
178 pass
179 clock = date.strftime(clock_format)
7500d90b 180
606def7e
BR
181 # Pull extended retweet text
182 try:
18df6e7f
O
183 text = 'RT @' + t['retweeted_status']['user']['screen_name'] + ': ' +\
184 t['retweeted_status']['text']
606def7e
BR
185 except:
186 pass
187
18df6e7f 188 # Unescape HTML character
606def7e
BR
189 text = unescape(text)
190
7500d90b
VNM
191 # Get expanded url
192 try:
193 expanded_url = []
194 url = []
195 urls = t['entities']['urls']
196 for u in urls:
197 expanded_url.append(u['expanded_url'])
198 url.append(u['url'])
199 except:
200 expanded_url = None
201 url = None
202
203 # Get media
204 try:
205 media_url = []
206 media = t['entities']['media']
207 for m in media:
208 media_url.append(m['media_url'])
209 except:
210 media_url = None
211
212 # Filter and ignore
213 screen_name = '@' + screen_name
214 if fil and screen_name not in fil:
215 return
216 if ig and screen_name in ig:
217 return
218
219 # Get rainbow id
220 res = db.tweet_to_rainbow_query(tid)
221 if not res:
222 db.tweet_store(tid)
223 res = db.tweet_to_rainbow_query(tid)
224 rid = res[0].rainbow_id
225
226 # Format info
8c7ae594 227 name = cycle_color(name)
d6cc4c67 228 nick = color_func(c['TWEET']['nick'])(screen_name)
0d9977c7
O
229 clock = clock
230 id = str(rid)
8c7ae594 231 fav = ''
7500d90b 232 if favorited:
8c7ae594
O
233 fav = color_func(c['TWEET']['favorited'])(u'\u2605')
234
7500d90b
VNM
235 tweet = text.split()
236 # Replace url
237 if expanded_url:
238 for index in range(len(expanded_url)):
c3bab4ef 239 tweet = lmap(
8141aca6 240 lambda x: expanded_url[index]
241 if x == url[index]
242 else x,
7500d90b
VNM
243 tweet)
244 # Highlight RT
c3bab4ef 245 tweet = lmap(
8141aca6 246 lambda x: color_func(c['TWEET']['rt'])(x)
247 if x == 'RT'
248 else x,
c075e6dc 249 tweet)
7500d90b 250 # Highlight screen_name
c3bab4ef 251 tweet = lmap(lambda x: cycle_color(x) if x[0] == '@' else x, tweet)
7500d90b 252 # Highlight link
c3bab4ef 253 tweet = lmap(
8141aca6 254 lambda x: color_func(c['TWEET']['link'])(x)
255 if x[0:4] == 'http'
256 else x,
c075e6dc 257 tweet)
4c025026 258 # Highlight hashtag
259 tweet = lmap(
8141aca6 260 lambda x: color_func(c['TWEET']['hashtag'])(x)
261 if x.startswith('#')
262 else x,
4c025026 263 tweet)
59262e95 264 # Highlight keyword
7500d90b 265 tweet = ' '.join(tweet)
59262e95 266 if keyword:
a8c5fce4 267 roj = re.search(keyword, tweet, re.IGNORECASE)
59262e95
O
268 if roj:
269 occur = roj.group()
270 ary = tweet.split(occur)
b13c6a0d 271 delimiter = color_func(c['TWEET']['keyword'])(occur)
272 tweet = delimiter.join(ary)
7500d90b 273
8c7ae594
O
274 # Load config formater
275 try:
276 formater = c['FORMAT']['TWEET']['DISPLAY']
0d9977c7
O
277 formater = name.join(formater.split("#name"))
278 formater = nick.join(formater.split("#nick"))
279 formater = fav.join(formater.split("#fav"))
280 formater = tweet.join(formater.split("#tweet"))
281 # Change clock word
282 word = [w for w in formater.split() if '#clock' in w][0]
8141aca6 283 delimiter = color_func(c['TWEET']['clock'])(
284 clock.join(word.split('#clock')))
0d9977c7
O
285 formater = delimiter.join(formater.split(word))
286 # Change id word
287 word = [w for w in formater.split() if '#id' in w][0]
288 delimiter = color_func(c['TWEET']['id'])(id.join(word.split('#id')))
289 formater = delimiter.join(formater.split(word))
8c7ae594
O
290 except:
291 printNicely(red('Wrong format in config.'))
292 return
7500d90b 293
8c7ae594
O
294 # Draw
295 printNicely(formater)
7500d90b
VNM
296
297 # Display Image
fe9bb33b 298 if c['IMAGE_ON_TERM'] and media_url:
7500d90b 299 for mu in media_url:
17bc529d 300 try:
301 response = requests.get(mu)
77f1d210 302 image_to_display(BytesIO(response.content))
303 except Exception:
17bc529d 304 printNicely(red('Sorry, image link is broken'))
7500d90b
VNM
305
306
c37c04a9 307def print_message(m, check_semaphore=False):
7500d90b
VNM
308 """
309 Print direct message
310 """
c37c04a9
O
311
312 # Check the semaphore lock (stream process only)
313 if check_semaphore:
314 if db.semaphore_query_pause():
315 return
316 while db.semaphore_query_lock():
317 time.sleep(0.5)
318
319 # Retrieve message
7500d90b
VNM
320 sender_screen_name = '@' + m['sender_screen_name']
321 sender_name = m['sender']['name']
b2cde062 322 text = unescape(m['text'])
7500d90b
VNM
323 recipient_screen_name = '@' + m['recipient_screen_name']
324 recipient_name = m['recipient']['name']
325 mid = m['id']
326 date = parser.parse(m['created_at'])
327 date = date - datetime.timedelta(seconds=time.timezone)
8c7ae594
O
328 clock_format = '%Y/%m/%d %H:%M:%S'
329 try:
330 clock_format = c['FORMAT']['MESSAGE']['CLOCK_FORMAT']
331 except:
332 pass
333 clock = date.strftime(clock_format)
7500d90b
VNM
334
335 # Get rainbow id
336 res = db.message_to_rainbow_query(mid)
337 if not res:
338 db.message_store(mid)
339 res = db.message_to_rainbow_query(mid)
340 rid = res[0].rainbow_id
341
6fa09c14 342 # Draw
8c7ae594
O
343 sender_name = cycle_color(sender_name)
344 sender_nick = color_func(c['MESSAGE']['sender'])(sender_screen_name)
345 recipient_name = cycle_color(recipient_name)
0d9977c7
O
346 recipient_nick = color_func(
347 c['MESSAGE']['recipient'])(recipient_screen_name)
8c7ae594 348 to = color_func(c['MESSAGE']['to'])('>>>')
0d9977c7
O
349 clock = clock
350 id = str(rid)
8c7ae594 351
c3bab4ef 352 text = ''.join(lmap(lambda x: x + ' ' if x == '\n' else x, text))
7500d90b 353
8c7ae594
O
354 # Load config formater
355 try:
356 formater = c['FORMAT']['MESSAGE']['DISPLAY']
0d9977c7
O
357 formater = sender_name.join(formater.split("#sender_name"))
358 formater = sender_nick.join(formater.split("#sender_nick"))
359 formater = to.join(formater.split("#to"))
360 formater = recipient_name.join(formater.split("#recipient_name"))
361 formater = recipient_nick.join(formater.split("#recipient_nick"))
362 formater = text.join(formater.split("#message"))
363 # Change clock word
364 word = [w for w in formater.split() if '#clock' in w][0]
8141aca6 365 delimiter = color_func(c['MESSAGE']['clock'])(
366 clock.join(word.split('#clock')))
0d9977c7
O
367 formater = delimiter.join(formater.split(word))
368 # Change id word
369 word = [w for w in formater.split() if '#id' in w][0]
370 delimiter = color_func(c['MESSAGE']['id'])(id.join(word.split('#id')))
371 formater = delimiter.join(formater.split(word))
8c7ae594
O
372 except:
373 printNicely(red('Wrong format in config.'))
374 return
375
376 # Draw
377 printNicely(formater)
7500d90b
VNM
378
379
fe9bb33b 380def show_profile(u):
7500d90b
VNM
381 """
382 Show a profile
383 """
384 # Retrieve info
385 name = u['name']
386 screen_name = u['screen_name']
387 description = u['description']
388 profile_image_url = u['profile_image_url']
389 location = u['location']
390 url = u['url']
391 created_at = u['created_at']
392 statuses_count = u['statuses_count']
393 friends_count = u['friends_count']
394 followers_count = u['followers_count']
6fa09c14 395
7500d90b 396 # Create content
c075e6dc
O
397 statuses_count = color_func(
398 c['PROFILE']['statuses_count'])(
399 str(statuses_count) +
400 ' tweets')
401 friends_count = color_func(
402 c['PROFILE']['friends_count'])(
403 str(friends_count) +
404 ' following')
405 followers_count = color_func(
406 c['PROFILE']['followers_count'])(
407 str(followers_count) +
408 ' followers')
7500d90b 409 count = statuses_count + ' ' + friends_count + ' ' + followers_count
c075e6dc
O
410 user = cycle_color(
411 name) + color_func(c['PROFILE']['nick'])(' @' + screen_name + ' : ') + count
412 profile_image_raw_url = 'Profile photo: ' + \
413 color_func(c['PROFILE']['profile_image_url'])(profile_image_url)
7500d90b 414 description = ''.join(
c3bab4ef 415 lmap(lambda x: x + ' ' * 4 if x == '\n' else x, description))
632c6fa5
O
416 description = color_func(c['PROFILE']['description'])(description)
417 location = 'Location : ' + color_func(c['PROFILE']['location'])(location)
418 url = 'URL : ' + (color_func(c['PROFILE']['url'])(url) if url else '')
7500d90b
VNM
419 date = parser.parse(created_at)
420 date = date - datetime.timedelta(seconds=time.timezone)
421 clock = date.strftime('%Y/%m/%d %H:%M:%S')
632c6fa5 422 clock = 'Join at ' + color_func(c['PROFILE']['clock'])(clock)
6fa09c14 423
7500d90b
VNM
424 # Format
425 line1 = u"{u:>{uw}}".format(
426 u=user,
427 uw=len(user) + 2,
428 )
429 line2 = u"{p:>{pw}}".format(
430 p=profile_image_raw_url,
431 pw=len(profile_image_raw_url) + 4,
432 )
433 line3 = u"{d:>{dw}}".format(
434 d=description,
435 dw=len(description) + 4,
436 )
437 line4 = u"{l:>{lw}}".format(
438 l=location,
439 lw=len(location) + 4,
440 )
441 line5 = u"{u:>{uw}}".format(
442 u=url,
443 uw=len(url) + 4,
444 )
445 line6 = u"{c:>{cw}}".format(
446 c=clock,
447 cw=len(clock) + 4,
448 )
6fa09c14 449
7500d90b
VNM
450 # Display
451 printNicely('')
452 printNicely(line1)
fe9bb33b 453 if c['IMAGE_ON_TERM']:
17bc529d 454 try:
455 response = requests.get(profile_image_url)
c37c04a9 456 image_to_display(BytesIO(response.content))
17bc529d 457 except:
458 pass
7500d90b
VNM
459 else:
460 printNicely(line2)
461 for line in [line3, line4, line5, line6]:
462 printNicely(line)
463 printNicely('')
464
465
466def print_trends(trends):
467 """
468 Display topics
469 """
632c6fa5 470 for topic in trends[:c['TREND_MAX']]:
7500d90b
VNM
471 name = topic['name']
472 url = topic['url']
8394e34b 473 line = cycle_color(name) + ': ' + color_func(c['TREND']['url'])(url)
7500d90b
VNM
474 printNicely(line)
475 printNicely('')
2d341029
O
476
477
478def print_list(group):
479 """
480 Display a list
481 """
482 for g in group:
483 # Format
422dd385 484 name = g['full_name']
2d341029
O
485 name = color_func(c['GROUP']['name'])(name + ' : ')
486 member = str(g['member_count'])
422dd385 487 member = color_func(c['GROUP']['member'])(member + ' member')
2d341029 488 subscriber = str(g['subscriber_count'])
422dd385
O
489 subscriber = color_func(
490 c['GROUP']['subscriber'])(
491 subscriber +
492 ' subscriber')
2d341029
O
493 description = g['description'].strip()
494 description = color_func(c['GROUP']['description'])(description)
495 mode = g['mode']
422dd385 496 mode = color_func(c['GROUP']['mode'])('Type: ' + mode)
2d341029
O
497 created_at = g['created_at']
498 date = parser.parse(created_at)
499 date = date - datetime.timedelta(seconds=time.timezone)
500 clock = date.strftime('%Y/%m/%d %H:%M:%S')
501 clock = 'Created at ' + color_func(c['GROUP']['clock'])(clock)
502
2d341029 503 # Create lines
422dd385
O
504 line1 = ' ' * 2 + name + member + ' ' + subscriber
505 line2 = ' ' * 4 + description
506 line3 = ' ' * 4 + mode
507 line4 = ' ' * 4 + clock
2d341029
O
508
509 # Display
510 printNicely('')
511 printNicely(line1)
512 printNicely(line2)
513 printNicely(line3)
514 printNicely(line4)
515
516 printNicely('')
59262e95
O
517
518
8141aca6 519def show_calendar(month, date, rel):
520 """
521 Show the calendar in rainbow mode
522 """
523 month = random_rainbow(month)
524 date = ' '.join([cycle_color(i) for i in date.split(' ')])
525 today = str(int(os.popen('date +\'%d\'').read().strip()))
526 # Display
527 printNicely(month)
528 printNicely(date)
529 for line in rel:
530 ary = line.split(' ')
531 ary = lmap(
532 lambda x: color_func(c['CAL']['today'])(x)
533 if x == today
534 else color_func(c['CAL']['days'])(x),
535 ary)
536 printNicely(' '.join(ary))
537
538
59262e95 539# Start the color cycle
b2cde062 540start_cycle()