diff --git a/README.md b/README.md index c52d05a7b..706d43db0 100644 --- a/README.md +++ b/README.md @@ -389,8 +389,35 @@ if err != nil { defer f.Close() scanner := chess.NewScanner(f) -for scanner.Scan() { - game := scanner.Next() +// Read all games +for scanner.HasNext() { + game, err := scanner.ParseNext() + if err != nil { + log.Fatal("Failed to parse game: %v", err) + } + fmt.Println(game.GetTagPair("Site")) + // Output &{Site https://lichess.org/8jb5kiqw} +} +``` + +#### Scan PGN expanding all variations + +To expand every variation into a distinct Game: + +```go +f, err := os.Open("lichess_db_standard_rated_2013-01.pgn") +if err != nil { + panic(err) +} +defer f.Close() + +scanner := chess.NewScanner(f, chess.WithExpandVariations()) +// Read all games +for scanner.HasNext() { + game, err := scanner.ParseNext() + if err != nil { + log.Fatal("Failed to parse game: %v", err) + } fmt.Println(game.GetTagPair("Site")) // Output &{Site https://lichess.org/8jb5kiqw} } diff --git a/game.go b/game.go index a22a85fd3..887335856 100644 --- a/game.go +++ b/game.go @@ -897,3 +897,67 @@ func (g *Game) updatePosition(move *Move) { move.position = newPos } } + +// Split takes a Game with a main line and 0 or more variations and returns a +// slice of Games (one for each variation), each containing exactly only a main +// line and 0 variations +func (g *Game) Split() []*Game { + // Collect all move paths starting from the root's children + var paths [][]*Move + for _, m := range g.rootMove.children { + for _, p := range collectPaths(m) { + paths = append(paths, p) + } + } + + // Build a Game for each path + var games []*Game + for _, path := range paths { + newG := g.buildOneGameFromPath(path) + games = append(games, newG) + } + + return games +} + +// collectPaths returns all paths from the given move to each leaf node. +// Each path is represented as a slice of *Move, starting with the given node +// and ending with a leaf (a move with no children). +func collectPaths(node *Move) [][]*Move { + if node == nil { + return nil + } + // If leaf, return a single path containing this node + if len(node.children) == 0 { + return [][]*Move{{node}} + } + // Otherwise, collect paths from each child and prepend this node + var paths [][]*Move + for _, c := range node.children { + childPaths := collectPaths(c) + for _, p := range childPaths { + path := append([]*Move{node}, p...) + paths = append(paths, path) + } + } + return paths +} + +func (g *Game) buildOneGameFromPath(path []*Move) *Game { + rootMove := &Move{position: g.rootMove.position.copy()} + cur := rootMove + + for _, m := range path { + child := m.Clone() + child.parent = cur + + cur.children = []*Move{child} + cur = child + } + + newG := g.Clone() + newG.rootMove = rootMove + newG.currentMove = cur + + return newG +} diff --git a/game_test.go b/game_test.go index c281f46d0..8a14d98df 100644 --- a/game_test.go +++ b/game_test.go @@ -1220,3 +1220,57 @@ func FuzzTestPushNotationMove(f *testing.F) { _ = game.PushNotationMove(move, notation, nil) }) } + +func validateSplit(t *testing.T, origPgn string, expectedLastLines []string) { + reader := strings.NewReader(origPgn) + scanner := NewScanner(reader) + game, err := scanner.ParseNext() + if err != nil { + t.Fatalf("fail to parse game: %s", err.Error()) + } + + if game == nil { + t.Fatalf("game is nil") + } + + splitGames := game.Split() + if len(expectedLastLines) != len(splitGames) { + t.Fatalf("expected %v split games but got %v", len(expectedLastLines), + len(splitGames)) + } + + for idx, g := range game.Split() { + lines := strings.Split(g.String(), "\n") + if len(lines) == 0 { + t.Fatalf("split game %v output blank", idx) + } + + lastLine := lines[len(lines)-1] + if lastLine != expectedLastLines[idx] { + t.Errorf("game output not correct\n\tExpected:'%v'\n\tGot: '%v'\n", + expectedLastLines[idx], lastLine) + } + } +} + +func TestGameSplitVar(t *testing.T) { + expectedLastLines := []string{ + "1. e4 e5 2. Nf3 Nc6 3. d4 exd4 4. Nxd4 *", + "1. e4 e5 2. Nc3 Nf6 3. f4 *", + "1. e4 d6 2. d4 Nf6 3. Nc3 e5 4. dxe5 dxe5 5. Qxd8+ Kxd8 *", + "1. e4 d6 2. d4 Nf6 3. Nc3 e5 4. Nf3 Nbd7 *", + "1. e3 e5 *", + } + + pgn := mustParsePGN("fixtures/pgns/variations.pgn") + validateSplit(t, pgn, expectedLastLines) +} + +func TestGameSplitNoVar(t *testing.T) { + expectedLastLines := []string{ + "1. e4 e5 2. Nf3 Nc6 *", + } + + pgn := "[Event \"SomeEvent\"]\n1. e4 e5 2. Nf3 Nc6\n\n" + validateSplit(t, pgn, expectedLastLines) +} diff --git a/move.go b/move.go index 8d2aab9df..9938f034d 100644 --- a/move.go +++ b/move.go @@ -147,3 +147,30 @@ func (m *Move) Ply() int { // After the move, it's Black's turn, so the move was by White return (moveNumber)*2 + 0 } + +// Clone returns a deep copy of a move. +// +// Per-field exceptions: +// +// parent: not copied; the clone'd move has no parent +// children: not copied; the clone'd move has no children +func (m *Move) Clone() *Move { + ret := &Move{} + ret.parent = nil + ret.position = m.position.copy() + ret.nag = m.nag + ret.comments = m.comments + ret.children = make([]*Move, 0) + ret.number = m.number + ret.tags = m.tags + ret.s1 = m.s1 + ret.s2 = m.s2 + ret.promo = m.promo + + ret.command = make(map[string]string) + for k, v := range m.command { + ret.command[k] = v + } + + return ret +} diff --git a/move_test.go b/move_test.go index 7ff47ecee..f76006449 100644 --- a/move_test.go +++ b/move_test.go @@ -439,3 +439,68 @@ func moveIsValid(pos *Position, m *Move, useTags bool) bool { } return false } + +func assertMovesAreEqual(t *testing.T, m1, m2 *Move) { + if m1.parent != m2.parent { + t.Fatalf("cloned mv %s parent is not the same", m1) + } + if m1.position.String() != m2.position.String() { + t.Fatalf("cloned mv %s position is not the same", m1) + } + if m1.nag != m2.nag { + t.Fatalf("cloned mv %s nag is not the same", m1) + } + if m1.comments != m2.comments { + t.Fatalf("cloned mv %s comments is not the same", m1) + } + if m1.number != m2.number { + t.Fatalf("cloned mv %s number is not the same", m1) + } + if m1.tags != m2.tags { + t.Fatalf("cloned mv %s tags is not the same", m1) + } + if m1.s1 != m2.s1 { + t.Fatalf("cloned mv %s s1 is not the same", m1) + } + if m1.s2 != m2.s2 { + t.Fatalf("cloned mv %s s2 is not the same", m1) + } + if m1.promo != m2.promo { + t.Fatalf("cloned mv %s s2 is not the same", m1) + } + + if len(m1.command) != len(m2.command) { + t.Fatalf("cloned mv %s len(command) is not the same", m1) + } else { + for k, v1 := range m1.command { + v2, ok := m2.command[k] + if !ok || v2 != v1 { + t.Fatalf("cloned mv %s command[%v] is not the same", m1, k) + } + } + } + if len(m1.children) != len(m2.children) { + t.Fatalf("cloned mv %s len(command) is not the same", m1) + } else { + for idx, c1 := range m1.children { + c2 := m2.children[idx] + assertMovesAreEqual(t, c1, c2) + } + } +} + +func TestMoveClone(t *testing.T) { + for _, mt := range validMoves { + mt.m.position = mt.pos + clonedM1 := mt.m.Clone() + assertMovesAreEqual(t, mt.m, clonedM1) + clonedM1.SetCommand("foo", "bar") + clonedM2 := clonedM1.Clone() + assertMovesAreEqual(t, clonedM1, clonedM2) + clonedM1.SetCommand("foo", "bar modified") + fooVal, ok := clonedM2.GetCommand("foo") + if !ok || fooVal != "bar" { + t.Fatalf("cloned mv %s is not a deep copy", clonedM2) + } + } +} diff --git a/pgn_test.go b/pgn_test.go index 8583de9a3..e12a6e1e6 100644 --- a/pgn_test.go +++ b/pgn_test.go @@ -91,22 +91,10 @@ func TestGamesFromPGN(t *testing.T) { for idx, test := range validPGNs { reader := strings.NewReader(test.PGN) scanner := NewScanner(reader) - scannedGame, err := scanner.ScanGame() + game, err := scanner.ParseNext() if err != nil { - t.Fatalf("fail to scan game from valid pgn %d: %s", idx, err.Error()) + t.Fatalf("fail to parse game from valid pgn %d: %s", idx, err.Error()) } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("fail to tokenize game from valid pgn %d: %s", idx, err.Error()) - } - - parser := NewParser(tokens) - game, err := parser.Parse() - if err != nil { - t.Fatalf("fail to read games from valid pgn %d: %s", idx, err.Error()) - } - if game == nil { t.Fatalf("game is nil") } @@ -118,22 +106,10 @@ func TestGameWithVariations(t *testing.T) { reader := strings.NewReader(pgn) scanner := NewScanner(reader) - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("fail to scan game from valid pgn: %s", err.Error()) - } - - tokens, err := TokenizeGame(scannedGame) + game, err := scanner.ParseNext() if err != nil { - t.Fatalf("fail to tokenize game from valid pgn: %s", err.Error()) + t.Fatalf("fail to parse game from pgn: %s", err.Error()) } - - parser := NewParser(tokens) - game, err := parser.Parse() - if err != nil { - t.Fatalf("fail to read games from valid pgn: %s", err.Error()) - } - if game == nil { t.Fatalf("game is nil") } @@ -161,18 +137,7 @@ func TestSingleGameFromPGN(t *testing.T) { scanner := NewScanner(reader) - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("fail to scan game from valid pgn: %s", err.Error()) - } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("fail to tokenize game from valid pgn: %s", err.Error()) - } - - parser := NewParser(tokens) - game, err := parser.Parse() + game, err := scanner.ParseNext() if err != nil { t.Fatalf("fail to read games from valid pgn: %s", err.Error()) } @@ -331,18 +296,7 @@ func TestCompleteGame(t *testing.T) { reader := strings.NewReader(pgn) scanner := NewScanner(reader) - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("fail to scan game from valid pgn: %s", err.Error()) - } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("fail to tokenize game from valid pgn: %s", err.Error()) - } - - parser := NewParser(tokens) - game, err := parser.Parse() + game, err := scanner.ParseNext() if err != nil { t.Fatalf("fail to read games from valid pgn: %s", err.Error()) } @@ -413,18 +367,7 @@ func TestLichessMultipleCommand(t *testing.T) { scanner := NewScanner(file) // Test first game - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("Failed to read first game: %v", err) - } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("Failed to tokenize first game: %v", err) - } - - parser := NewParser(tokens) - game, err := parser.Parse() + game, err := scanner.ParseNext() if err != nil { t.Fatalf("fail to read games from valid pgn: %s", err.Error()) } @@ -475,18 +418,7 @@ func TestParseMoveWithNAGAndComment(t *testing.T) { 1. e4 $1 {Good move} e5 {Solid} $2 2. Nf3 $3 {Another comment} Nc6 $4 {Yet another}` scanner := NewScanner(strings.NewReader(pgn)) - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("fail to scan game: %v", err) - } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("fail to tokenize game: %v", err) - } - - parser := NewParser(tokens) - game, err := parser.Parse() + game, err := scanner.ParseNext() if err != nil { t.Fatalf("fail to parse game: %v", err) } @@ -522,18 +454,7 @@ func TestVariationMoveNumbers(t *testing.T) { 1. e4 e5 2. Nf3 Nc6 3. Bb5 (3. Bc4 Nf6 4. d3) a6 4. Ba4 Nf6 5. O-O Be7 1-0` scanner := NewScanner(strings.NewReader(pgn)) - scannedGame, err := scanner.ScanGame() - if err != nil { - t.Fatalf("fail to scan game: %v", err) - } - - tokens, err := TokenizeGame(scannedGame) - if err != nil { - t.Fatalf("fail to tokenize game: %v", err) - } - - parser := NewParser(tokens) - game, err := parser.Parse() + game, err := scanner.ParseNext() if err != nil { t.Fatalf("fail to parse game: %v", err) } diff --git a/scanner.go b/scanner.go index 67d5da412..4bec5751a 100644 --- a/scanner.go +++ b/scanner.go @@ -11,15 +11,12 @@ Example usage: // Read all games for scanner.HasNext() { - game, err := scanner.ScanGame() + game, err := scanner.ParseNext() if err != nil { - log.Fatal(err) + log.Fatalf("Failed to parse game: %v", err) } // Process game } - - // Tokenize a specific game - tokens, err := TokenizeGame(game) */ package chess @@ -71,9 +68,26 @@ func TokenizeGame(game *GameScanned) ([]Token, error) { // It supports streaming processing of multiple games and proper handling // of PGN syntax. type Scanner struct { - scanner *bufio.Scanner - nextGame *GameScanned // Buffer for peeked game - lastError error // Store last error + scanner *bufio.Scanner + nextGame *GameScanned // Buffer for peeked game + lastError error // Store last error + opts ScannerOpts + nextParsedGames []*Game // only valid when ExpandVariations==true +} + +type ScannerOption func(*Scanner) + +// WithExpandVariations() instructs the scanner to expand all variations in +// a single GameScanned into multiple Game instances (1 per variation) rather +// than a single Game instance. +func WithExpandVariations() ScannerOption { + return func(s *Scanner) { + s.opts.ExpandVariations = true + } +} + +type ScannerOpts struct { + ExpandVariations bool // default false } // NewScanner creates a new PGN scanner that reads from the provided reader. @@ -83,10 +97,20 @@ type Scanner struct { // Example: // // scanner := NewScanner(strings.NewReader(pgnText)) -func NewScanner(r io.Reader) *Scanner { +func NewScanner(r io.Reader, opts ...ScannerOption) *Scanner { s := bufio.NewScanner(r) s.Split(splitPGNGames) - return &Scanner{scanner: s} + ret := &Scanner{ + scanner: s, + nextParsedGames: make([]*Game, 0), + } + + // apply all the options + for _, opt := range opts { + opt(ret) + } + + return ret } // ScanGame reads and returns the next game from the source. @@ -125,12 +149,12 @@ func (s *Scanner) ScanGame() (*GameScanned, error) { // Example: // // for scanner.HasNext() { -// game, err := scanner.ScanGame() -// // Process game +// scangame, err := scanner.ScanGame() +// // Process scangame // } func (s *Scanner) HasNext() bool { // If we already have a buffered game, return true - if s.nextGame != nil { + if s.nextGame != nil || len(s.nextParsedGames) > 0 { return true } @@ -146,6 +170,45 @@ func (s *Scanner) HasNext() bool { return false } +// ParseNext is a convenience wrapper combining the functionality of +// ScanGame(), TokenizeGame(), NewParser(), and Parse() enabling +// callers to simplify iterating over each Game within a pgn file. +// +// Example: +// +// for scanner.HasNext() { +// game, err := scanner.ParseNext() +// // Process game +// } +func (s *Scanner) ParseNext() (*Game, error) { + if len(s.nextParsedGames) > 0 { + ret := s.nextParsedGames[0] + s.nextParsedGames = s.nextParsedGames[1:] + return ret, nil + } + + scannedGame, err := s.ScanGame() + if err != nil { + return nil, err + } + tokens, err := TokenizeGame(scannedGame) + if err != nil { + return nil, err + } + parser := NewParser(tokens) + game, err := parser.Parse() + if err != nil { + return nil, err + } + if !s.opts.ExpandVariations { + return game, nil + } // else + + parsedGames := game.Split() + s.nextParsedGames = parsedGames[1:] + return parsedGames[0], nil +} + // Split function for bufio.Scanner to split PGN games. func splitPGNGames(data []byte, atEOF bool) (int, []byte, error) { // Skip leading whitespace diff --git a/scanner_test.go b/scanner_test.go index f92398685..774068e81 100644 --- a/scanner_test.go +++ b/scanner_test.go @@ -5,6 +5,7 @@ import ( "io" "os" "path/filepath" + "strings" "testing" ) @@ -230,3 +231,63 @@ func TestHasNextDoesntConsume(t *testing.T) { t.Error("First game has no tokens after multiple HasNext calls") } } + +func validateExpand(t *testing.T, scanner *Scanner, expectedLastLines []string) { + count := 0 + for scanner.HasNext() { + game, err := scanner.ParseNext() + if err != nil { + t.Fatalf("fail to parse game: %s", err.Error()) + } + + if game == nil { + t.Fatalf("game is nil") + } + if count >= len(expectedLastLines) { + t.Fatalf("expected %v games but found at least %v", + len(expectedLastLines), count+1) + } + lines := strings.Split(game.String(), "\n") + if len(lines) == 0 { + t.Fatalf("split game %v output blank", count+1) + } + + lastLine := lines[len(lines)-1] + if lastLine != expectedLastLines[count] { + t.Errorf("game output not correct\n\tExpected:'%v'\n\tGot: '%v'\n", + expectedLastLines[count], lastLine) + } + count++ + } + + if count != len(expectedLastLines) { + t.Fatalf("expected %v games but found only %v", + len(expectedLastLines), count) + } +} + +func TestScannerExpand(t *testing.T) { + expectedLastLines := []string{ + "1. e4 e5 2. Nf3 Nc6 3. d4 exd4 4. Nxd4 *", + "1. e4 e5 2. Nc3 Nf6 3. f4 *", + "1. e4 d6 2. d4 Nf6 3. Nc3 e5 4. dxe5 dxe5 5. Qxd8+ Kxd8 *", + "1. e4 d6 2. d4 Nf6 3. Nc3 e5 4. Nf3 Nbd7 *", + "1. e3 e5 *", + } + + pgn := mustParsePGN("fixtures/pgns/variations.pgn") + reader := strings.NewReader(pgn) + scanner := NewScanner(reader, WithExpandVariations()) + validateExpand(t, scanner, expectedLastLines) +} + +func TestScannerNoExpand(t *testing.T) { + expectedLastLines := []string{ + "1. e4 (1. e3 e5) 1... e5 (1... d6 2. d4 Nf6 3. Nc3 e5 4. dxe5 (4. Nf3 Nbd7) 4... dxe5 5. Qxd8+ Kxd8) 2. Nf3 (2. Nc3 Nf6 3. f4) 2... Nc6 3. d4 exd4 4. Nxd4 *", + } + + pgn := mustParsePGN("fixtures/pgns/variations.pgn") + reader := strings.NewReader(pgn) + scanner := NewScanner(reader) + validateExpand(t, scanner, expectedLastLines) +}