berkerpeksag / astor

Python AST read/write
https://pypi.org/project/astor/
BSD 3-Clause "New" or "Revised" License
810 stars 102 forks source link

code_to_ast can't extract methods in classes #71

Open miau opened 7 years ago

miau commented 7 years ago

Passing a method to code_to_ast causes the following error.

>>> import astor
>>> from requests.auth import HTTPDigestAuth
>>> node = astor.code_to_ast(HTTPDigestAuth.build_digest_header)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "C:\Users\miau\Documents\repo\astor\astor\file_util.py", line 102, in __call__
    return cache[key]
KeyError: ('C:\\ProgramData\\Anaconda3\\lib\\site-packages\\requests\\auth.py', 128)

It seems that astor doesn't look into ast.ClassDef.

https://github.com/berkerpeksag/astor/blob/0a02019977ced2a8450583179d3537b9e4978d7e/astor/file_util.py#L99

pmaupin commented 7 years ago

Yeah, that function's pretty simple right now -- it only works on top-level functions.

Having said that, you can have functions inside classes inside functions inside functions inside classes inside...

Honestly, I don't even remember my use-case for writing that, but I suspect it had to do with testing the code writer.

We'd certainly accept a patch for a more-functional version of this, but I don't have a use-case at the moment.

Thanks, Pat

zmitchell commented 6 years ago

I wrote a version of the CodeToAst class for myself that recurses down the AST. Here's basically how it works.

I make a list of AST nodes, block_types, that have a body field (ast.If, ast.For, etc), and thus could contain a function. I also make a list of AST nodes, func_types, that are function definitions (ast.FunctionDef or ast.AsyncFunctionDef).

Then I have a method that descends the AST (note self.asts is equivalent to cache, and self.filename is equivalent to fname):

def _find_funcs(self, parent_ast):
    for item in parent_ast.body:
        if type(item) in self.func_types:
            self.asts[(self.filename, item.lineno)] = item
            self._find_funcs(item)
        elif type(item) in self.block_types:
            self._find_funcs(item)

I haven't tried to port this back to astor yet, but I could do that without much trouble if it sounds like what you're looking for.

berkerpeksag commented 6 years ago

@zmitchell that seems like a reasonable approach to me, thank you. I'd be happy to review if you submit a PR.

zmitchell commented 6 years ago

Sounds good, I'll put it together  👍