Can you configure python-mode, elpy-mode to apply font-locking to imported library functions?

Question

This question is likely an expansion / duplicate of this question: how to highlight imported functions in python code?

My .py file looks like this:

from math import floor

(max(1, 8))

(floor(3.7))

Screenshot for reference:

max has font-lock applied to it. While floor has no font-lock applied.

Questions:

1. How can I configure my setup (python-mode, elpy-mode) so that emacs recognises 'floor' as a function that font-lock should apply to.
2. If there are no answers for #1, can you recommend a package which can help?

Answer 1

I think this might be close to what you want:

The next-import function leverages font lock to dynamically add these keywords. Then, we add a hook function that runs that.

(defun next-import (limit)
  (when (re-search-forward "from \\(\\w+\\) import \\(\\w+\\)" limit t)
    (font-lock-add-keywords
     nil
     `((,(regexp-quote (match-string 2)) 0 font-lock-keyword-face)))))

(defun my-import-keywords ()
  (font-lock-add-keywords
   nil
   '((next-import (nil nil t)))
   t))

(add-hook 'python-mode-hook 'my-import-keywords)

One limitation of this is that it only font-locks from imports that have a single imported function. With some work you could probably match a comma separated list, and add each one.