Method _extract_class

kernelbench/bridge.py:324–347

Extract a class definition (including body) from source.

(self, class_name: str)

'''

    def _extract_class(self, class_name: str) -> Optional[str]:
        """Extract a class definition (including body) from source."""
        lines = self.source_code.split("\n")
        collecting = False
        class_lines: List[str] = []

        for line in lines:
            if re.match(rf"^class\s+{class_name}\s*\(", line):
                collecting = True
                class_lines.append(line)
                continue
            if collecting:
                # End of class: non-empty, non-indented, top-level construct
                if line.strip() and not line[0].isspace() and not line.strip().startswith("#"):
                    if re.match(r"^(class |def |@)", line):
                        break
                class_lines.append(line)

        if class_lines:
            # Trim trailing blank lines
            while class_lines and not class_lines[-1].strip():
                class_lines.pop()
            return "\n".join(class_lines)
        return None


# ---------------------------------------------------------------------------