AST get_source_segment throws the following error, TypeError: object of type 'Module' has no len()

Question

I am trying to code a python AST that can unparse the if statement node. Here is my code -

import ast

tree = ast.parse("""

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def age_change(self):
        if(self.age>20):
            return self.age + 5
        else:
            return self.age

        """)

for node in ast.walk(tree):
    if(isinstance(node, ast.If)):
        print(ast.get_source_segment(tree, node))

However, when I do that, I get the following error -

Traceback (most recent call last):
  File "C:\Users\thoma\AppData\Roaming\JetBrains\PyCharmCE2020.3\scratches\scratch_21.py", line 20, in <module>
    print(ast.get_source_segment(tree, node))
  File "C:\Users\thoma\anaconda3\envs\torch_1\lib\ast.py", line 354, in get_source_segment
    lines = _splitlines_no_ff(source)
  File "C:\Users\thoma\anaconda3\envs\torch_1\lib\ast.py", line 307, in _splitlines_no_ff
    while idx < len(source):
TypeError: object of type 'Module' has no len()

I got the documentation for get_source_segment from here - https://docs.python.org/3/library/ast.html and therefore am not sure why I am getting this error.

Answer 1

@jasonharper is correct. The first parameter should be a string having the source code. Here is the correct code -

import ast
code = """

class Person:
    def __init__(self, name, age):
        self.name = name
        self.age = age

    def age_change(self):
        if(self.age>20):
            return self.age + 5
        else:
            return self.age

        """
tree = ast.parse(code)

for node in ast.walk(tree):
    if(isinstance(node, ast.If)):
        print(ast.get_source_segment(code, node))