How can I get source code of imported libraries functions?

Answered

Good time of day!

I am developing a plugin for PyCharm that analyzes functions code and predicts inputs and outputs for it. For now I am only available to get methods that were written inside a project (by developer). I do it in a following way [I also suggest "suitable" variables] (I will not attach all the code, because there are lots of not needed stuff):

@Override
protected void addCompletions(@NotNull CompletionParameters parameters,
@NotNull ProcessingContext context,
@NotNull CompletionResultSet result) {
.....
CaretModel caret = parameters.getEditor().getCaretModel();
PsiElement currentElem = parameters.getOriginalFile().findElementAt(caret.getOffset());
....
PyCallExpression callExpression = PsiTreeUtil.getParentOfType(currentElem, PyCallExpression.class);
PyExpression callee = Objects.requireNonNull(callExpression).getCallee();
.....
PyGotoDeclarationHandler handler = new PyGotoDeclarationHandler();
PsiElement pyFunction = handler.getGotoDeclarationTarget(callee, parameters.getEditor());
.....
String functionCode = pyFunction.getText()

The question is: how can I access functions that are written in libraries?
Let's say, that user somewhere inside code imported tensorflow library in a next way:

import tensorflow as tf

and later he used:

tf.exp(5)

My aim is to get the code of the tf.exp() function that is:

def exp(self, context=None):
"""Returns e ** self."""

if context is None:
context = getcontext()

# exp(NaN) = NaN
ans = self._check_nans(context=context)
if ans:
return ans
....
return ans

I tried playing with scopes and FindUsagesProvider but didn't succeed. Another my try was connected to Documentation providing, but for such case I need an original PsiElement which I don't know ho to get. Also, I thought about using Psi Files, but didn't find a way out.

How can I get all "def " functions? Not obligatory with PyGotoDeclarationHandler. I hope that there exists some universal solution :)

Thank's in advance!

1
4 comments
Official comment

Hi!

Check an example here, call.multiResolveCallee(resolveContext) returns possible callees for the call.

It is not recommended to call getText on the elements that are not yet parsed (opened in the editor) since it leads to file parsing and makes your plugin much slower.

Thank you very much for the answer! It works!

So, I have a question regarding getText: the point is that my plugin needs the whole function body (the model analyses the text of function, and after that predicts possible output types). What is the appropriate way of getting the whole text of a function in such a case? 

Is it only time issued? (I am asking this because model inference also takes some time, maybe getText time is not so crucial in such context...).

I am very interested to know the best way of implementing it.
Thank you one more time)

Edited: inference time is faster :(

0

Mostly time consuming, yes.

It could be achieved with custom index: during indexing process your indexer receives files to be indexed, checks if there are any functions that it is interested in, runs model and saves results. Then, during code completion (or other user/IDE activity), these results are taken from the index and used for further processing.

Examples:

  • `com.intellij.util.indexing.FileBasedIndexExtension`,
  • `com.jetbrains.python.psi.stubs.PyModuleNameIndex`,
  • `com.jetbrains.python.psi.stubs.PySetuptoolsNamespaceIndex`,
  • implementations of `com.intellij.util.indexing.IndexExtension#getIndexer`

Feel free to post questions here.

PS Is your plugin going to support predefined set of functions? or it should be ready for any function in a user's codebase?

0

Thanks for the examples! I will take a look at them.

As for your question: for any function in a user's codebase. If I am correct, there are 2 types in common:

1) Functions that were written by the user of PyCharm, and look like:

def simple_fun(a, b):
return a+b

2) And the second type as in the example that I already mentioned:

tf.exp(5)

We want to handle at least these cases. If there are other cases that we forgot - it would be nice to handle them too.

0

Please sign in to leave a comment.