mentions.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. import re
  2. from bs4 import BeautifulSoup, NavigableString
  3. from django.contrib.auth import get_user_model
  4. from django.utils import six
  5. SUPPORTED_TAGS = ('h1', 'h2', 'h3', 'h4', 'h5', 'h6', 'p')
  6. USERNAME_RE = re.compile(r'@[0-9a-z]+', re.IGNORECASE)
  7. MENTIONS_LIMIT = 24
  8. def add_mentions(request, result):
  9. if '@' not in result['parsed_text']:
  10. return
  11. mentions_dict = {}
  12. soup = BeautifulSoup(result['parsed_text'], 'html5lib')
  13. for tagname in SUPPORTED_TAGS:
  14. if tagname in result['parsed_text']:
  15. for element in soup.find_all(tagname):
  16. add_mentions_to_element(request, element, mentions_dict)
  17. result['parsed_text'] = six.text_type(soup.body)[6:-7].strip()
  18. result['mentions'] = filter(bool, mentions_dict.values())
  19. def add_mentions_to_element(request, element, mentions_dict):
  20. for item in element.contents:
  21. if item.name:
  22. if item.name != 'a':
  23. add_mentions_to_element(request, item, mentions_dict)
  24. elif '@' in item.string:
  25. parse_string(request, item, mentions_dict)
  26. def parse_string(request, element, mentions_dict):
  27. def replace_mentions(matchobj):
  28. if len(mentions_dict) >= MENTIONS_LIMIT:
  29. return matchobj.group(0)
  30. username = matchobj.group(0)[1:].strip().lower()
  31. if username not in mentions_dict:
  32. if username == request.user.slug:
  33. mentions_dict[username] = request.user
  34. else:
  35. User = get_user_model()
  36. try:
  37. mentions_dict[username] = User.objects.get(slug=username)
  38. except User.DoesNotExist:
  39. mentions_dict[username] = None
  40. if mentions_dict[username]:
  41. user = mentions_dict[username]
  42. return u'<a href="{}">@{}</a>'.format(user.get_absolute_url(), user.username)
  43. else:
  44. # we've failed to resolve user for username
  45. return matchobj.group(0)
  46. replaced_string = USERNAME_RE.sub(replace_mentions, element.string)
  47. element.replace_with(BeautifulSoup(replaced_string, 'html.parser'))