Skip to content

Commit b5cfea3

Browse files
authored
fix: update modules in sanitize.py
1 parent 3fbcbb6 commit b5cfea3

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

bigcodebench/sanitize.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,8 @@
66
from pqdm.processes import pqdm
77

88
from tqdm import tqdm
9-
from tree_sitter import Node
10-
from tree_sitter_languages import get_parser
9+
import tree_sitter_python
10+
from tree_sitter import Language, Node, Parser
1111

1212
from bigcodebench.data import (
1313
get_bigcodebench,
@@ -111,7 +111,7 @@ def has_return_statement(node: Node) -> bool:
111111
def extract_target_code_or_empty(code: str, entrypoint: Optional[str] = None) -> str:
112112
code = code_extract(code.strip())
113113
code_bytes = bytes(code, "utf8")
114-
parser = get_parser("python")
114+
parser = Parser(Language(tree_sitter_python.language()))
115115
tree = parser.parse(code_bytes)
116116
class_names = set()
117117
function_names = set()
@@ -299,4 +299,4 @@ def main():
299299

300300

301301
if __name__ == "__main__":
302-
main()
302+
main()

0 commit comments

Comments
 (0)