Automate the Boring Stuff with Go, Ch. 9: Excel Files

go get github.com/xuri/excelize/v2

// read-workbook opens an Excel file, lists its sheets, prints the first sheet as a
// table and shows how to read one cell by its name.
//
//	go run ./ch9/projects/read-workbook ./ch9/projects/read-workbook/example.xlsx
//	go run ./ch9/projects/read-workbook -cell C3 ./ch9/projects/read-workbook/example.xlsx
package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	"strings"
	"text/tabwriter"

	"github.com/xuri/excelize/v2"
)

// Table returns the rows of sheet, padded so every row has as many cells as the widest.
func Table(f *excelize.File, sheet string) ([][]string, error) {
	rows, err := f.GetRows(sheet)
	if err != nil {
		return nil, fmt.Errorf("rows of %s: %w", sheet, err)
	}
	width := 0
	for _, r := range rows {
		width = max(width, len(r))
	}
	for i := range rows {
		for len(rows[i]) < width {
			rows[i] = append(rows[i], "")
		}
	}
	return rows, nil
}

func run(args []string) error {
	flags := flag.NewFlagSet("read-workbook", flag.ContinueOnError)
	cell := flags.String("cell", "B2", "one cell to print by name")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if flags.NArg() != 1 {
		return errors.New("usage: read-workbook [-cell B2] <file.xlsx>")
	}
	f, err := excelize.OpenFile(flags.Arg(0))
	if err != nil {
		return fmt.Errorf("open: %w", err)
	}
	defer f.Close()

	sheets := f.GetSheetList()
	fmt.Println("sheets:", strings.Join(sheets, ", "))
	first := sheets[0]

	rows, err := Table(f, first)
	if err != nil {
		return err
	}
	tw := tabwriter.NewWriter(os.Stdout, 0, 4, 2, ' ', 0)
	for _, r := range rows {
		fmt.Fprintln(tw, strings.Join(r, "\t"))
	}
	if err := tw.Flush(); err != nil {
		return err
	}

	v, err := f.GetCellValue(first, *cell)
	if err != nil {
		return fmt.Errorf("cell %s: %w", *cell, err)
	}
	fmt.Printf("\n%s!%s = %q\n", first, *cell, v)
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// census-by-county reads a census workbook (one row per census tract) and totals the
// population and the tract count per county, then writes the result as JSON so other
// programs can use it.
//
//	go run ./ch9/projects/census-by-county ./ch9/projects/census-by-county/census.xlsx
package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"slices"
	"strconv"

	"github.com/xuri/excelize/v2"
)

// Tract is one row of the workbook.
type Tract struct {
	State  string
	County string
	Pop    int
}

// CountyStats is what we know about one county after reading every tract.
type CountyStats struct {
	Tracts int `json:"tracts"`
	Pop    int `json:"pop"`
}

// ReadTracts loads the first sheet. It expects the columns CensusTract, State, County,
// POP2010 in that order, with a header row.
func ReadTracts(path string) ([]Tract, error) {
	f, err := excelize.OpenFile(path)
	if err != nil {
		return nil, fmt.Errorf("open: %w", err)
	}
	defer f.Close()
	rows, err := f.GetRows(f.GetSheetList()[0])
	if err != nil {
		return nil, err
	}
	var tracts []Tract
	for i, row := range rows {
		if i == 0 || len(row) < 4 {
			continue
		}
		pop, err := strconv.Atoi(row[3])
		if err != nil {
			return nil, fmt.Errorf("row %d: population %q: %w", i+1, row[3], err)
		}
		tracts = append(tracts, Tract{State: row[1], County: row[2], Pop: pop})
	}
	return tracts, nil
}

// Aggregate returns state -> county -> stats.
func Aggregate(tracts []Tract) map[string]map[string]CountyStats {
	out := map[string]map[string]CountyStats{}
	for _, t := range tracts {
		if out[t.State] == nil {
			out[t.State] = map[string]CountyStats{}
		}
		s := out[t.State][t.County]
		s.Tracts++
		s.Pop += t.Pop
		out[t.State][t.County] = s
	}
	return out
}

func run(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: census-by-county <census.xlsx>")
	}
	tracts, err := ReadTracts(args[0])
	if err != nil {
		return err
	}
	stats := Aggregate(tracts)

	for _, state := range slices.Sorted(keys(stats)) {
		total := 0
		for _, c := range stats[state] {
			total += c.Pop
		}
		fmt.Printf("%s: %d people in %d counties\n", state, total, len(stats[state]))
		for _, county := range slices.Sorted(keys(stats[state])) {
			c := stats[state][county]
			fmt.Printf("  %-12s %6d people, %d tracts\n", county, c.Pop, c.Tracts)
		}
	}

	data, err := json.MarshalIndent(stats, "", "  ")
	if err != nil {
		return err
	}
	if err := os.WriteFile("census_summary.json", data, 0o644); err != nil {
		return fmt.Errorf("write summary: %w", err)
	}
	fmt.Println("wrote census_summary.json")
	return nil
}

func keys[V any](m map[string]V) func(func(string) bool) {
	return func(yield func(string) bool) {
		for k := range m {
			if !yield(k) {
				return
			}
		}
	}
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// write-workbook creates an Excel file from scratch: a named sheet, a bold header,
// a few rows, sensible column widths and a SUM formula.
//
//	go run ./ch9/projects/write-workbook
package main

import (
	"fmt"
	"os"

	"github.com/xuri/excelize/v2"
)

// Item is one line of inventory.
type Item struct {
	Name     string
	Quantity int
	Price    float64
}

// BuildInventory fills a new workbook with the items and returns it unsaved.
func BuildInventory(items []Item) (*excelize.File, error) {
	f := excelize.NewFile()
	const sheet = "Inventory"
	if err := f.SetSheetName("Sheet1", sheet); err != nil {
		return nil, err
	}
	header := []any{"Item", "Quantity", "Unit price", "Value"}
	if err := f.SetSheetRow(sheet, "A1", &header); err != nil {
		return nil, err
	}
	bold, err := f.NewStyle(&excelize.Style{Font: &excelize.Font{Bold: true}})
	if err != nil {
		return nil, err
	}
	if err := f.SetCellStyle(sheet, "A1", "D1", bold); err != nil {
		return nil, err
	}
	for i, it := range items {
		row := i + 2
		values := []any{it.Name, it.Quantity, it.Price}
		if err := f.SetSheetRow(sheet, fmt.Sprintf("A%d", row), &values); err != nil {
			return nil, err
		}
		if err := f.SetCellFormula(sheet, fmt.Sprintf("D%d", row), fmt.Sprintf("B%d*C%d", row, row)); err != nil {
			return nil, err
		}
	}
	last := len(items) + 1
	if err := f.SetCellValue(sheet, fmt.Sprintf("A%d", last+1), "Total"); err != nil {
		return nil, err
	}
	if err := f.SetCellFormula(sheet, fmt.Sprintf("D%d", last+1), fmt.Sprintf("SUM(D2:D%d)", last)); err != nil {
		return nil, err
	}
	if err := f.SetColWidth(sheet, "A", "A", 18); err != nil {
		return nil, err
	}
	if err := f.SetColWidth(sheet, "B", "D", 12); err != nil {
		return nil, err
	}
	return f, nil
}

func run() error {
	items := []Item{
		{"Keyboard", 12, 39.50},
		{"Mouse", 30, 14.99},
		{"Monitor", 5, 219.00},
		{"USB-C cable", 80, 6.25},
	}
	f, err := BuildInventory(items)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := f.SaveAs("inventory.xlsx"); err != nil {
		return fmt.Errorf("save: %w", err)
	}
	fmt.Println("wrote inventory.xlsx")
	return nil
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// update-produce-prices fixes wrong unit prices in a sales workbook and recomputes the
// TOTAL column for every row it touched. The corrections live in one map, so the next
// time prices change you edit three lines instead of three hundred cells.
//
//	go run ./ch9/projects/update-produce-prices ./ch9/projects/update-produce-prices/produceSales.xlsx
package main

import (
	"errors"
	"fmt"
	"os"
	"strconv"

	"github.com/xuri/excelize/v2"
)

var corrections = map[string]float64{
	"Garlic": 3.07,
	"Celery": 1.19,
	"Lemon":  1.27,
}

// UpdatePrices rewrites COST PER POUND (column B) and TOTAL (column D) for every row
// whose PRODUCE (column A) has a corrected price. It returns how many rows changed.
func UpdatePrices(f *excelize.File, sheet string, prices map[string]float64) (int, error) {
	rows, err := f.GetRows(sheet)
	if err != nil {
		return 0, err
	}
	changed := 0
	for i, row := range rows {
		if i == 0 || len(row) < 3 {
			continue
		}
		price, ok := prices[row[0]]
		if !ok {
			continue
		}
		pounds, err := strconv.ParseFloat(row[2], 64)
		if err != nil {
			return changed, fmt.Errorf("row %d: pounds %q: %w", i+1, row[2], err)
		}
		n := i + 1
		if err := f.SetCellValue(sheet, fmt.Sprintf("B%d", n), price); err != nil {
			return changed, err
		}
		if err := f.SetCellValue(sheet, fmt.Sprintf("D%d", n), round2(price*pounds)); err != nil {
			return changed, err
		}
		changed++
	}
	return changed, nil
}

func round2(x float64) float64 { return float64(int(x*100+0.5)) / 100 }

func run(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: update-produce-prices <produceSales.xlsx>")
	}
	f, err := excelize.OpenFile(args[0])
	if err != nil {
		return fmt.Errorf("open: %w", err)
	}
	defer f.Close()
	sheet := f.GetSheetList()[0]
	n, err := UpdatePrices(f, sheet, corrections)
	if err != nil {
		return err
	}
	out := "updatedProduceSales.xlsx"
	if err := f.SaveAs(out); err != nil {
		return fmt.Errorf("save: %w", err)
	}
	fmt.Printf("updated %d row(s), wrote %s\n", n, out)
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// multiplication-table writes an N by N multiplication table to a workbook, with the
// first row and column in bold.
//
//	go run ./ch9/projects/multiplication-table -n 12
package main

import (
	"flag"
	"fmt"
	"os"

	"github.com/xuri/excelize/v2"
)

// Fill writes the table for 1..n into sheet, leaving A1 empty like a paper table.
func Fill(f *excelize.File, sheet string, n int) error {
	bold, err := f.NewStyle(&excelize.Style{Font: &excelize.Font{Bold: true}})
	if err != nil {
		return err
	}
	for i := 1; i <= n; i++ {
		rowHeader, err := excelize.CoordinatesToCellName(1, i+1)
		if err != nil {
			return err
		}
		colHeader, err := excelize.CoordinatesToCellName(i+1, 1)
		if err != nil {
			return err
		}
		for _, cell := range []string{rowHeader, colHeader} {
			if err := f.SetCellValue(sheet, cell, i); err != nil {
				return err
			}
			if err := f.SetCellStyle(sheet, cell, cell, bold); err != nil {
				return err
			}
		}
		for j := 1; j <= n; j++ {
			cell, err := excelize.CoordinatesToCellName(j+1, i+1)
			if err != nil {
				return err
			}
			if err := f.SetCellValue(sheet, cell, i*j); err != nil {
				return err
			}
		}
	}
	return nil
}

func run(args []string) error {
	flags := flag.NewFlagSet("multiplication-table", flag.ContinueOnError)
	n := flags.Int("n", 10, "table size")
	if err := flags.Parse(args); err != nil {
		return err
	}
	if *n < 1 || *n > 1000 {
		return fmt.Errorf("n must be between 1 and 1000, got %d", *n)
	}
	f := excelize.NewFile()
	defer f.Close()
	if err := Fill(f, "Sheet1", *n); err != nil {
		return err
	}
	out := fmt.Sprintf("multiplication_table_%d.xlsx", *n)
	if err := f.SaveAs(out); err != nil {
		return fmt.Errorf("save: %w", err)
	}
	fmt.Println("wrote", out)
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// txt-to-xlsx takes text files and writes one workbook where each file is a column and
// each line is a cell.
//
//	go run ./ch9/projects/txt-to-xlsx ch9/projects/txt-to-xlsx/first.txt ch9/projects/txt-to-xlsx/second.txt ch9/projects/txt-to-xlsx/third.txt
package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/xuri/excelize/v2"
)

// ReadLines returns the lines of a file without their line endings.
func ReadLines(path string) ([]string, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	var lines []string
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		lines = append(lines, sc.Text())
	}
	return lines, sc.Err()
}

// ColumnsToWorkbook puts each slice of lines into its own column of a new workbook, with
// the given names as a header row.
func ColumnsToWorkbook(names []string, columns [][]string) (*excelize.File, error) {
	f := excelize.NewFile()
	for c := range columns {
		col, err := excelize.ColumnNumberToName(c + 1)
		if err != nil {
			return nil, err
		}
		values := append([]any{names[c]}, toAny(columns[c])...)
		if err := f.SetSheetCol("Sheet1", col+"1", &values); err != nil {
			return nil, err
		}
	}
	return f, nil
}

func toAny(s []string) []any {
	out := make([]any, len(s))
	for i, v := range s {
		out[i] = v
	}
	return out
}

func run(args []string) error {
	if len(args) == 0 {
		return errors.New("usage: txt-to-xlsx <file.txt> [more.txt ...]")
	}
	var names []string
	var columns [][]string
	for _, path := range args {
		lines, err := ReadLines(path)
		if err != nil {
			return fmt.Errorf("%s: %w", path, err)
		}
		names = append(names, strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)))
		columns = append(columns, lines)
	}
	f, err := ColumnsToWorkbook(names, columns)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := f.SaveAs("combined.xlsx"); err != nil {
		return fmt.Errorf("save: %w", err)
	}
	fmt.Printf("wrote combined.xlsx with %d column(s)\n", len(columns))
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}

// xlsx-to-txt is the reverse of txt-to-xlsx: every column of the first sheet becomes a
// text file named after the column letter, one cell per line.
//
//	go run ./ch9/projects/xlsx-to-txt ./ch9/projects/xlsx-to-txt/table.xlsx
package main

import (
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"github.com/xuri/excelize/v2"
)

// Columns reads the first sheet column by column, trimming trailing empty cells.
func Columns(f *excelize.File) ([][]string, error) {
	cols, err := f.GetCols(f.GetSheetList()[0])
	if err != nil {
		return nil, err
	}
	for i := range cols {
		for len(cols[i]) > 0 && cols[i][len(cols[i])-1] == "" {
			cols[i] = cols[i][:len(cols[i])-1]
		}
	}
	return cols, nil
}

// WriteColumns writes each column to dir/column_<letter>.txt and returns the paths.
func WriteColumns(dir string, cols [][]string) ([]string, error) {
	var written []string
	for i, col := range cols {
		letter, err := excelize.ColumnNumberToName(i + 1)
		if err != nil {
			return written, err
		}
		path := filepath.Join(dir, "column_"+letter+".txt")
		if err := os.WriteFile(path, []byte(strings.Join(col, "\n")+"\n"), 0o644); err != nil {
			return written, err
		}
		written = append(written, path)
	}
	return written, nil
}

func run(args []string) error {
	if len(args) != 1 {
		return errors.New("usage: xlsx-to-txt <file.xlsx>")
	}
	f, err := excelize.OpenFile(args[0])
	if err != nil {
		return fmt.Errorf("open: %w", err)
	}
	defer f.Close()
	cols, err := Columns(f)
	if err != nil {
		return err
	}
	paths, err := WriteColumns(".", cols)
	if err != nil {
		return err
	}
	for _, p := range paths {
		fmt.Println("wrote", p)
	}
	return nil
}

func main() {
	if err := run(os.Args[1:]); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}