mentions.py 2.1 KB

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