Extract the CUDA_SRC string from a Python kernel file. Handles: - CUDA_SRC = r\"\"\"...\"\"\" - CUDA_SRC = \"\"\"...\"\"\" - CUDA_SRC = r'''...''' - CUDA_SRC = '''...''' Returns the raw CUDA C++ source string, or None if not found.
(source: str)
| 97 | # --------------------------------------------------------------------------- |
| 98 | |
| 99 | def extract_cuda_source(source: str) -> Optional[str]: |
| 100 | """ |
| 101 | Extract the CUDA_SRC string from a Python kernel file. |
| 102 | |
| 103 | Handles: |
| 104 | - CUDA_SRC = r\"\"\"...\"\"\" |
| 105 | - CUDA_SRC = \"\"\"...\"\"\" |
| 106 | - CUDA_SRC = r'''...''' |
| 107 | - CUDA_SRC = '''...''' |
| 108 | |
| 109 | Returns the raw CUDA C++ source string, or None if not found. |
| 110 | """ |
| 111 | # Try AST-based extraction first (most robust) |
| 112 | try: |
| 113 | tree = ast.parse(source) |
| 114 | for node in ast.walk(tree): |
| 115 | if isinstance(node, ast.Assign): |
| 116 | for target in node.targets: |
| 117 | if isinstance(target, ast.Name) and target.id == "CUDA_SRC": |
| 118 | if isinstance(node.value, ast.Constant) and isinstance( |
| 119 | node.value.value, str |
| 120 | ): |
| 121 | return node.value.value |
| 122 | # Python 3.7 compat: ast.Str |
| 123 | if hasattr(ast, "Str") and isinstance(node.value, ast.Str): |
| 124 | return node.value.s |
| 125 | except SyntaxError: |
| 126 | pass |
| 127 | |
| 128 | # Fallback: regex-based extraction for triple-quoted strings |
| 129 | # Matches: CUDA_SRC = r"""...""" or CUDA_SRC = """...""" |
| 130 | for quote in ('"""', "'''"): |
| 131 | pattern = rf'CUDA_SRC\s*=\s*r?{re.escape(quote)}(.*?){re.escape(quote)}' |
| 132 | match = re.search(pattern, source, re.DOTALL) |
| 133 | if match: |
| 134 | return match.group(1) |
| 135 | |
| 136 | return None |
| 137 | |
| 138 | |
| 139 | def extract_function_name_from_compile(source: str) -> Optional[str]: |
no outgoing calls
no test coverage detected