File size: 4,080 Bytes
97e363b |
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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
import dataclasses
from typing import Optional
from .artifacts import AbstractArtifact
import ast
from .artifacts import FunctionArtifact
@dataclasses.dataclass(frozen=True)
class Program:
"""A parsed Python program."""
# `preface` is everything from the beginning of the code till the first
# function is found.
preface: str
artifacts: list[AbstractArtifact]
def __str__(self) -> str:
program = f'{self.preface}\n' if self.preface else ''
program += '\n'.join([str(f) for f in self.artifacts])
return program
def find_artifact_index(self, artifact_name: str) -> int:
"""Returns the index of input function name."""
artifact_names = [a.name for a in self.artifacts]
count = artifact_names.count(artifact_name)
if count == 0:
raise ValueError(
f'artifact {artifact_name} does not exist in program:\n{str(self)}'
)
if count > 1:
raise ValueError(
f'artifact {artifact_name} exists more than once in program:\n'
f'{str(self)}'
)
index = artifact_names.index(artifact_name)
return index
def get_artifact(self, artifact_name: str) -> AbstractArtifact:
index = self.find_artifact_index(artifact_name)
return self.artifacts[index]
# TODO: Do this for various types of artifacts (only for functions rn)
class ProgramVisitor(ast.NodeVisitor):
"""Parses code to collect all required information to produce a `Program`.
Note that we do not store function decorators.
"""
def __init__(self, sourcecode: str):
self._codelines: list[str] = sourcecode.splitlines()
self._preface: str = ''
self._artifacts: list[AbstractArtifact] = []
self._current_artifact: Optional[str] = None
def visit_FunctionDef(self,
node: ast.FunctionDef):
"""Collects all information about the function being parsed."""
if node.col_offset == 0: # We only care about first level functions.
self._current_function = node.name
if not self._artifacts:
self._preface = '\n'.join(self._codelines[:node.lineno - 1])
function_end_line = node.end_lineno
body_start_line = node.body[0].lineno - 1
# Extract the docstring.
docstring = None
if isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value,ast.Str):
docstring = f' """{ast.literal_eval(ast.unparse(node.body[0]))}"""'
if len(node.body) > 1:
body_start_line = node.body[1].lineno - 1
else:
body_start_line = function_end_line
self._artifacts.append(FunctionArtifact(
name=node.name,
args=ast.unparse(node.args),
return_type=ast.unparse(node.returns) if node.returns else None,
docstring=docstring,
body='\n'.join(self._codelines[body_start_line:function_end_line]),
))
self.generic_visit(node)
def return_program(self) -> Program:
return Program(preface=self._preface, artifacts=self._artifacts)
def text_to_program(text: str) -> Program:
"""Returns Program object by parsing input text using Python AST."""
# We assume that the program is composed of some preface (e.g. imports,
# classes, assignments, ...) followed by a sequence of functions.
#Often happens that it returns a codeblock so remove it
if text.startswith("```python"):
text = text[9:]
if text.endswith("```"):
text = text[:-3]
tree = ast.parse(text)
visitor = ProgramVisitor(text)
visitor.visit(tree)
return visitor.return_program()
def text_to_artifact(text: str) -> AbstractArtifact:
"""Returns Function object by parsing input text using Python AST."""
program = text_to_program(text)
if len(program.artifacts) != 1:
raise ValueError(f'Only one artifact expected, got {len(program.artifacts)}'
f':\n{program.functions}')
return program.artifacts[0] |