FilesExpand file tree

main

/

source_code_parsing_utils.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

52 lines (45 loc) · 1.82 KB

main

/

source_code_parsing_utils.py

Copy path

Top

File metadata and controls

• Code
• Blame

52 lines (45 loc) · 1.82 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

import ast

import importlib

import inspect

import textwrap

class ReturnNameVisitor(ast.NodeVisitor):

"""Thanks to ChatGPT for pairing."""

def __init__(self):

self.return_names = []

def visit_Return(self, node):

# Check if the return value is a tuple.

if isinstance(node.value, ast.Tuple):

for elt in node.value.elts:

if isinstance(elt, ast.Name):

self.return_names.append(elt.id)

else:

try:

self.return_names.append(ast.unparse(elt))

except Exception:

self.return_names.append(str(elt))

else:

if isinstance(node.value, ast.Name):

self.return_names.append(node.value.id)

else:

try:

self.return_names.append(ast.unparse(node.value))

except Exception:

self.return_names.append(str(node.value))

self.generic_visit(node)

def _determine_parent_module(self, cls):

from diffusers import DiffusionPipeline

from diffusers.models.modeling_utils import ModelMixin

if issubclass(cls, DiffusionPipeline):

return "pipelines"

elif issubclass(cls, ModelMixin):

return "models"

else:

raise NotImplementedError

def get_ast_tree(self, cls, attribute_name="encode_prompt"):

parent_module_name = self._determine_parent_module(cls)

main_module = importlib.import_module(f"diffusers.{parent_module_name}")

current_cls_module = getattr(main_module, cls.__name__)

source_code = inspect.getsource(getattr(current_cls_module, attribute_name))

source_code = textwrap.dedent(source_code)

tree = ast.parse(source_code)

return tree