68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
|
from typing import List
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Segment:
|
||
|
|
segment_id: str
|
||
|
|
start_line: int
|
||
|
|
end_line: int
|
||
|
|
start_byte: int
|
||
|
|
end_byte: int
|
||
|
|
content: str
|
||
|
|
|
||
|
|
|
||
|
|
class SourceSegmenter:
|
||
|
|
def __init__(self, code: str):
|
||
|
|
self.code = code
|
||
|
|
self.lines = code.split('\n')
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def segment(code: str) -> List[Segment]:
|
||
|
|
segmenter = SourceSegmenter(code)
|
||
|
|
return segmenter._find_segments()
|
||
|
|
|
||
|
|
def _find_segments(self) -> List[Segment]:
|
||
|
|
boundaries = self._find_boundaries()
|
||
|
|
segments = []
|
||
|
|
|
||
|
|
for i, (start, end) in enumerate(boundaries):
|
||
|
|
segment_id = f"seg_{i}"
|
||
|
|
start_byte = self._line_offset(start)
|
||
|
|
end_byte = self._line_offset(end)
|
||
|
|
content = '\n'.join(self.lines[start:end])
|
||
|
|
|
||
|
|
segments.append(Segment(
|
||
|
|
segment_id=segment_id,
|
||
|
|
start_line=start,
|
||
|
|
end_line=end,
|
||
|
|
start_byte=start_byte,
|
||
|
|
end_byte=end_byte,
|
||
|
|
content=content
|
||
|
|
))
|
||
|
|
|
||
|
|
return segments
|
||
|
|
|
||
|
|
def _line_offset(self, line_num: int) -> int:
|
||
|
|
total = 0
|
||
|
|
for i in range(min(line_num, len(self.lines))):
|
||
|
|
total += len(self.lines[i]) + 1
|
||
|
|
return total
|
||
|
|
|
||
|
|
def _find_boundaries(self) -> List[tuple]:
|
||
|
|
boundaries = []
|
||
|
|
start = 0
|
||
|
|
|
||
|
|
for i, line in enumerate(self.lines):
|
||
|
|
stripped = line.lstrip()
|
||
|
|
|
||
|
|
if stripped.startswith('def ') or stripped.startswith('class '):
|
||
|
|
if start < i and any(self.lines[start:i]):
|
||
|
|
boundaries.append((start, i))
|
||
|
|
start = i
|
||
|
|
|
||
|
|
if start < len(self.lines):
|
||
|
|
boundaries.append((start, len(self.lines)))
|
||
|
|
|
||
|
|
return boundaries if boundaries else [(0, len(self.lines))]
|