The function definition visiting method for tirx. Parameters ---------- self : Parser The visiting parser. node : doc.FunctionDef The doc AST function definition node.
(self: Parser, node: doc.FunctionDef)
| 596 | |
| 597 | @dispatch.register(token="tirx", type_name="FunctionDef") |
| 598 | def visit_function_def(self: Parser, node: doc.FunctionDef) -> None: |
| 599 | """The function definition visiting method for tirx. |
| 600 | |
| 601 | Parameters |
| 602 | ---------- |
| 603 | self : Parser |
| 604 | The visiting parser. |
| 605 | |
| 606 | node : doc.FunctionDef |
| 607 | The doc AST function definition node. |
| 608 | """ |
| 609 | supplied_annotation = self.function_annotations |
| 610 | func_annotation = supplied_annotation.get(node.name, {}) |
| 611 | privacy = find_decorator_annotation(node, "private", default=False) |
| 612 | s_tir = find_decorator_annotation(node, "s_tir", default=False) |
| 613 | persistent = find_decorator_annotation(node, "persistent", default=False) |
| 614 | self.function_annotations = None |
| 615 | with self.var_table.with_frame(): |
| 616 | prim_func_ctx = T.prim_func(is_private=privacy, s_tir=s_tir, persistent=persistent) |
| 617 | with prim_func_ctx: |
| 618 | T.func_name(node.name) |
| 619 | if node.returns is not None: |
| 620 | ret_type = self.eval_expr(node.returns) |
| 621 | if callable(ret_type): |
| 622 | ret_type = PrimType(ret_type().dtype) |
| 623 | T.func_ret(ret_type) |
| 624 | with self.with_dispatch_token("tirx"): |
| 625 | # TODO: handle different types of arguments: |
| 626 | # - vararg: arg | None |
| 627 | # - kwonlyargs: list[arg] |
| 628 | # - kw_defaults: list[expr | None] |
| 629 | # - kwarg: arg | None |
| 630 | # - defaults: list[expr] |
| 631 | # - posonlyargs: list[arg] |
| 632 | for arg in node.args.args: |
| 633 | if arg.annotation is None: |
| 634 | self.report_error(arg, "Type annotation required for function parameters.") |
| 635 | try: |
| 636 | ann = self.eval_expr(arg.annotation) |
| 637 | if callable(ann) and ann is not _constexpr_sentinel: |
| 638 | ann = ann() |
| 639 | except Exception: # pylint: disable=broad-except |
| 640 | ann = func_annotation.get(arg.arg, None) |
| 641 | if ann is None: |
| 642 | raise |
| 643 | if ann is _constexpr_sentinel: |
| 644 | # T.constexpr param: value was bound in extra_vars by |
| 645 | # TIRJit.specialize() and lives in an outer var_table |
| 646 | # frame; do not register a runtime PrimFunc param. |
| 647 | continue |
| 648 | param = T.arg(arg.arg, ann) |
| 649 | self.var_table.add(arg.arg, param) |
| 650 | self.visit_body(node.body) |
| 651 | self.function_annotations = supplied_annotation |
| 652 | |
| 653 | |
| 654 | @dispatch.register(token="tir.inline", type_name="FunctionDef") |
nothing calls this directly
no test coverage detected
searching dependent graphs…