diff --git a/main.go b/main.go index 92bd59e..39a87f7 100644 --- a/main.go +++ b/main.go @@ -252,26 +252,158 @@ func renderMDLine(line string, st *mdState) string { return indent + c(m[1]+" ", 33) + renderInline(m[2]) } if strings.HasPrefix(trimmed, "|") && strings.HasSuffix(trimmed, "|") { - isSep := true - for _, ch := range strings.ReplaceAll(strings.ReplaceAll(trimmed, "|", ""), ":", "") { - if ch != '-' && ch != ' ' { isSep = false; break } - } - if isSep { return c(trimmed, 2) } - cells := strings.Split(trimmed[1:len(trimmed)-1], "|") - var rcells []string - for _, cell := range cells { rcells = append(rcells, " "+renderInline(strings.TrimSpace(cell))+" ") } - return c("│", 2) + strings.Join(rcells, c("│", 2)) + c("│", 2) + return renderInline(line) } return renderInline(line) } +func isTableSep(cells []string) bool { + if len(cells) == 0 { return false } + for _, cell := range cells { + c := strings.ReplaceAll(strings.ReplaceAll(strings.TrimSpace(cell), "-", ""), ":", "") + if c != "" { return false } + } + return true +} + +func parseTableCells(line string) []string { + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, "|") { trimmed = trimmed[1:] } + if strings.HasSuffix(trimmed, "|") { trimmed = trimmed[:len(trimmed)-1] } + parts := strings.Split(trimmed, "|") + cells := make([]string, len(parts)) + for i, p := range parts { cells[i] = strings.TrimSpace(p) } + return cells +} + +func isTableLine(line string) bool { + trimmed := strings.TrimSpace(line) + if !strings.Contains(trimmed, "|") { return false } + return strings.HasPrefix(trimmed, "|") || strings.HasSuffix(trimmed, "|") +} + +func renderTable(lines []string) []string { + if len(lines) == 0 { return nil } + var rows [][]string + var headerRow []string + hasHeader := false + + for _, ln := range lines { + cells := parseTableCells(ln) + if isTableSep(cells) { + if len(rows) > 0 && !hasHeader { + headerRow = rows[len(rows)-1] + rows = rows[:len(rows)-1] + hasHeader = true + } + continue + } + rows = append(rows, cells) + } + + if !hasHeader && len(lines) < 2 { + var out []string + var st mdState + for _, ln := range lines { out = append(out, renderMDLine(ln, &st)) } + return out + } + + numCols := len(headerRow) + for _, r := range rows { + if len(r) > numCols { numCols = len(r) } + } + if numCols == 0 { return nil } + + if hasHeader { + for len(headerRow) < numCols { headerRow = append(headerRow, "") } + } + for i := range rows { + for len(rows[i]) < numCols { rows[i] = append(rows[i], "") } + } + + colWidths := make([]int, numCols) + for i := 0; i < numCols; i++ { + if hasHeader { + vl := visibleLen(headerRow[i]) + if vl > colWidths[i] { colWidths[i] = vl } + } + for _, r := range rows { + vl := visibleLen(renderInline(r[i])) + if vl > colWidths[i] { colWidths[i] = vl } + } + if colWidths[i] < 3 { colWidths[i] = 3 } + } + + var res []string + + var topParts []string + for _, w := range colWidths { topParts = append(topParts, strings.Repeat("─", w+2)) } + res = append(res, c("┌"+strings.Join(topParts, "┬")+"┐", 2)) + + if hasHeader { + var hCells []string + for i, h := range headerRow { + rh := c(h, 1, 36) + pad := strings.Repeat(" ", colWidths[i]-visibleLen(h)) + hCells = append(hCells, " "+rh+pad+" ") + } + res = append(res, c("│", 2)+strings.Join(hCells, c("│", 2))+c("│", 2)) + + var midParts []string + for _, w := range colWidths { midParts = append(midParts, strings.Repeat("─", w+2)) } + res = append(res, c("├"+strings.Join(midParts, "┼")+"┤", 2)) + } + + for _, r := range rows { + var rCells []string + for i, cell := range r { + rc := renderInline(cell) + pad := strings.Repeat(" ", colWidths[i]-visibleLen(rc)) + rCells = append(rCells, " "+rc+pad+" ") + } + res = append(res, c("│", 2)+strings.Join(rCells, c("│", 2))+c("│", 2)) + } + + var botParts []string + for _, w := range colWidths { botParts = append(botParts, strings.Repeat("─", w+2)) } + res = append(res, c("└"+strings.Join(botParts, "┴")+"┘", 2)) + + return res +} + func renderMD(text string) string { + if !COL { return text } lines := strings.Split(text, "\n") var out []string var st mdState + var tbl []string + + flushTable := func() { + if len(tbl) > 0 { + out = append(out, renderTable(tbl)...) + tbl = nil + } + } + for _, ln := range lines { + trimmed := strings.TrimSpace(ln) + if strings.HasPrefix(trimmed, "```") { + flushTable() + out = append(out, renderMDLine(ln, &st)) + continue + } + if st.inCode { + out = append(out, renderMDLine(ln, &st)) + continue + } + if isTableLine(ln) { + tbl = append(tbl, ln) + continue + } + flushTable() out = append(out, renderMDLine(ln, &st)) } + flushTable() return strings.Join(out, "\n") } @@ -420,8 +552,17 @@ func parseStream(r io.Reader) (Message, error) { var rh, ch bool var lineBuf string var mdSt mdState + var tblBuf []string tcs := map[int]*ToolCall{} var order []int + + flushTable := func() { + if len(tblBuf) > 0 { + for _, tln := range renderTable(tblBuf) { fmt.Println(tln) } + tblBuf = nil + } + } + sc := bufio.NewScanner(r) sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024) for sc.Scan() { @@ -449,7 +590,18 @@ func parseStream(r io.Reader) (Message, error) { if idx == -1 { break } curLine := lineBuf[:idx] lineBuf = lineBuf[idx+1:] - fmt.Println(renderMDLine(curLine, &mdSt)) + trimmed := strings.TrimSpace(curLine) + if strings.HasPrefix(trimmed, "```") { + flushTable() + fmt.Println(renderMDLine(curLine, &mdSt)) + } else if mdSt.inCode { + fmt.Println(renderMDLine(curLine, &mdSt)) + } else if isTableLine(curLine) { + tblBuf = append(tblBuf, curLine) + } else { + flushTable() + fmt.Println(renderMDLine(curLine, &mdSt)) + } } } for _, tc := range dl.ToolCalls { @@ -464,8 +616,14 @@ func parseStream(r io.Reader) (Message, error) { if tc.Function.Arguments != "" { t.Function.Arguments += tc.Function.Arguments } } } + flushTable() if lineBuf != "" { - fmt.Println(renderMDLine(lineBuf, &mdSt)) + if !mdSt.inCode && isTableLine(lineBuf) { + tblBuf = append(tblBuf, lineBuf) + flushTable() + } else { + fmt.Println(renderMDLine(lineBuf, &mdSt)) + } } if rh && !ch { fmt.Print("\n" + c("--- reasoning end ---", 36) + "\n") diff --git a/main_test.go b/main_test.go index 44fdfe7..572da42 100644 --- a/main_test.go +++ b/main_test.go @@ -1482,6 +1482,43 @@ func TestRenderMDBlocks(t *testing.T) { } } +func TestRenderTableFormatting(t *testing.T) { + COL = true + defer func() { COL = false }() + + raw := strings.Join([]string{ + "| Name | Role | Location |", + "| :--- | :---: | ---: |", + "| Alice | `Lead` | New York |", + "| Bob | Developer | London |", + }, "\n") + + rendered := renderMD(raw) + lines := strings.Split(rendered, "\n") + if len(lines) != 6 { + t.Fatalf("expected 6 table lines (top, header, mid, row1, row2, bot), got %d:\n%s", len(lines), rendered) + } + + if !strings.HasPrefix(lines[0], "\033[2m┌") || !strings.HasSuffix(lines[0], "┐\033[0m") { + t.Errorf("top border = %q", lines[0]) + } + if !strings.Contains(lines[1], "Name") || !strings.Contains(lines[1], "Role") || !strings.Contains(lines[1], "Location") { + t.Errorf("header row = %q", lines[1]) + } + if !strings.HasPrefix(lines[2], "\033[2m├") || !strings.HasSuffix(lines[2], "┤\033[0m") { + t.Errorf("mid border = %q", lines[2]) + } + if !strings.Contains(lines[3], "Alice") || !strings.Contains(lines[3], "New York") { + t.Errorf("row 1 = %q", lines[3]) + } + if !strings.Contains(lines[4], "Bob") || !strings.Contains(lines[4], "London") { + t.Errorf("row 2 = %q", lines[4]) + } + if !strings.HasPrefix(lines[5], "\033[2m└") || !strings.HasSuffix(lines[5], "┘\033[0m") { + t.Errorf("bot border = %q", lines[5]) + } +} + func TestRenderMDNoColorFallback(t *testing.T) { COL = false raw := "# Heading\n**bold** and `code`\n- list item" @@ -1492,3 +1529,4 @@ func TestRenderMDNoColorFallback(t *testing.T) { } +