From 83acc1e1afdfb3fc16bb235ed2bd9c85203c2ab3 Mon Sep 17 00:00:00 2001 From: Changkun Ou Date: Sun, 21 Dec 2025 12:57:06 +0100 Subject: [PATCH] all: simplify iterator API --- README.md | 14 ++++++++------ dataset_list.go | 13 +++++++------ dataset_test.go | 4 ++-- docs/dataset-apis.md | 4 ++-- docs/file-apis.md | 4 ++-- docs/parse-apis.md | 4 ++-- file_list.go | 6 +++--- file_test.go | 4 ++-- parse_classify_test.go | 2 +- parse_extract_test.go | 2 +- parse_get.go | 14 +++++++------- parse_get_test.go | 2 +- parse_list.go | 6 +++--- parse_parse_test.go | 4 ++-- util_test.go | 4 ++-- 15 files changed, 45 insertions(+), 42 deletions(-) diff --git a/README.md b/README.md index d0dad80..1543d30 100644 --- a/README.md +++ b/README.md @@ -194,12 +194,14 @@ result, err := c.GetParseResult( tensorlake.WithSSE(true), tensorlake.WithOnUpdate(func(name tensorlake.ParseEventName, r *tensorlake.ParseResult) { switch eventName { - case tensorlake.sseEventParseQueued: + case tensorlake.SSEEventParseQueued: fmt.Println("Job queued") - case tensorlake.sseEventParseUpdate: + case tensorlake.SSEEventParseUpdate: fmt.Printf("Progress: %d/%d pages\n", r.ParsedPagesCount, r.TotalPages) - case tensorlake.sseEventParseDone: + case tensorlake.SSEEventParseDone: fmt.Println("Complete!") + case tensorlake.SSEEventParseFailed: + fmt.Printf("Failed: %s\n", r.Error) } }), ) @@ -211,7 +213,7 @@ Easily iterate through paginated results: ```go // Iterate all files -for file, err := range c.IterFiles(ctx, 50, tensorlake.PaginationDirectionNext) { +for file, err := range c.IterFiles(ctx, 50) { if err != nil { panic(err) } @@ -219,7 +221,7 @@ for file, err := range c.IterFiles(ctx, 50, tensorlake.PaginationDirectionNext) } // Iterate all parse jobs -for job, err := range c.IterParseJobs(ctx, 50, tensorlake.PaginationDirectionNext) { +for job, err := range c.IterParseJobs(ctx, 50) { if err != nil { panic(err) } @@ -227,7 +229,7 @@ for job, err := range c.IterParseJobs(ctx, 50, tensorlake.PaginationDirectionNex } // Iterate all datasets -for dataset, err := range c.IterDatasets(ctx, 50, tensorlake.PaginationDirectionNext) { +for dataset, err := range c.IterDatasets(ctx, 50) { if err != nil { panic(err) } diff --git a/dataset_list.go b/dataset_list.go index c32507d..668b149 100644 --- a/dataset_list.go +++ b/dataset_list.go @@ -25,14 +25,14 @@ import ( ) // IterDatasets iterates over all datasets in the organization. -func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error] { +func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error] { return func(yield func(Dataset, error) bool) { cursor := "" for { listResp, err := c.ListDatasets(ctx, &ListDatasetsRequest{ Cursor: cursor, - Limit: limit, - Direction: direction, + Limit: batchSize, + Direction: PaginationDirectionNext, }) if err != nil { yield(Dataset{}, err) @@ -52,14 +52,15 @@ func (c *Client) IterDatasets(ctx context.Context, limit int, direction Paginati } // IterDatasetData iterates over all dataset data in the organization. -func (c *Client) IterDatasetData(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] { +func (c *Client) IterDatasetData(ctx context.Context, datasetId string, batchSize int) iter.Seq2[ParseResult, error] { return func(yield func(ParseResult, error) bool) { cursor := "" for { listResp, err := c.ListDatasetData(ctx, &ListDatasetDataRequest{ + DatasetId: datasetId, Cursor: cursor, - Limit: limit, - Direction: direction, + Limit: batchSize, + Direction: PaginationDirectionNext, }) if err != nil { yield(ParseResult{}, err) diff --git a/dataset_test.go b/dataset_test.go index 154ae3e..3eb8175 100644 --- a/dataset_test.go +++ b/dataset_test.go @@ -39,7 +39,7 @@ func TestDataset(t *testing.T) { // List datasets. datasets := []string{} - for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) { + for d, err := range c.IterDatasets(t.Context(), 1) { if err != nil { t.Fatalf("failed to list datasets: %v", err) } @@ -113,7 +113,7 @@ func TestDataset(t *testing.T) { // Check if the dataset is deleted. datasets = []string{} - for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) { + for d, err := range c.IterDatasets(t.Context(), 1) { if err != nil { t.Fatalf("failed to list datasets: %v", err) } diff --git a/docs/dataset-apis.md b/docs/dataset-apis.md index acadd71..4500f41 100644 --- a/docs/dataset-apis.md +++ b/docs/dataset-apis.md @@ -285,13 +285,13 @@ if response.HasMore { For convenience, use the `IterDatasets` method: ```go -func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error] +func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error] ``` **Example:** ```go -for dataset, err := range client.IterDatasets(context.Background(), 50, tensorlake.PaginationDirectionNext) { +for dataset, err := range client.IterDatasets(context.Background(), 50) { if err != nil { log.Fatal(err) } diff --git a/docs/file-apis.md b/docs/file-apis.md index aad3351..daf12e8 100644 --- a/docs/file-apis.md +++ b/docs/file-apis.md @@ -175,13 +175,13 @@ if response.HasMore { For convenience, use the `IterFiles` method to iterate through all files: ```go -func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error] +func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error] ``` **Example:** ```go -for file, err := range client.IterFiles(context.Background(), 50, tensorlake.PaginationDirectionNext) { +for file, err := range client.IterFiles(context.Background(), 50) { if err != nil { log.Fatal(err) } diff --git a/docs/parse-apis.md b/docs/parse-apis.md index 4f7fb07..176a55a 100644 --- a/docs/parse-apis.md +++ b/docs/parse-apis.md @@ -453,13 +453,13 @@ for _, job := range response.Items { ### Iterate All Parse Jobs ```go -func (c *Client) IterParseJobs(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] +func (c *Client) IterParseJobs(ctx context.Context, batchSize int) iter.Seq2[ParseResult, error] ``` **Example:** ```go -for job, err := range client.IterParseJobs(context.Background(), 50, tensorlake.PaginationDirectionNext) { +for job, err := range client.IterParseJobs(context.Background(), 50) { if err != nil { log.Fatal(err) } diff --git a/file_list.go b/file_list.go index 70bce75..c50f081 100644 --- a/file_list.go +++ b/file_list.go @@ -25,14 +25,14 @@ import ( ) // IterFiles iterates over all files in the project. -func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error] { +func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error] { return func(yield func(FileInfo, error) bool) { cursor := "" for { listResp, err := c.ListFiles(ctx, &ListFilesRequest{ Cursor: cursor, - Limit: limit, - Direction: direction, + Limit: batchSize, + Direction: PaginationDirectionNext, }) if err != nil { yield(FileInfo{}, err) diff --git a/file_test.go b/file_test.go index ab5eef0..7460b6d 100644 --- a/file_test.go +++ b/file_test.go @@ -80,7 +80,7 @@ func TestFileManagement(t *testing.T) { // List the files. Iterate through all the pages. files := []string{} - for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) { + for f, err := range c.IterFiles(t.Context(), 1) { if err != nil { t.Fatalf("failed to list files: %v", err) } @@ -124,7 +124,7 @@ func TestFileManagement(t *testing.T) { // Validate file is deleted. files = []string{} - for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) { + for f, err := range c.IterFiles(t.Context(), 1) { if err != nil { t.Fatalf("failed to list files: %v", err) } diff --git a/parse_classify_test.go b/parse_classify_test.go index 5ca2687..53432ab 100644 --- a/parse_classify_test.go +++ b/parse_classify_test.go @@ -101,7 +101,7 @@ func TestClassifyDocument(t *testing.T) { // Validate classify results. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) } diff --git a/parse_extract_test.go b/parse_extract_test.go index 0f7ec00..a5f8c3f 100644 --- a/parse_extract_test.go +++ b/parse_extract_test.go @@ -116,7 +116,7 @@ func TestExtractDocument(t *testing.T) { // Validate parse results. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) } diff --git a/parse_get.go b/parse_get.go index 5b55931..3257fb6 100644 --- a/parse_get.go +++ b/parse_get.go @@ -103,10 +103,10 @@ type ParseEventName string // The possible SSE events. // See also: https://github.com/tensorlakeai/tensorlake/blob/main/src/tensorlake/documentai/_parse.py#L499 const ( - sseEventParseQueued ParseEventName = "parse_queued" - sseEventParseUpdate ParseEventName = "parse_update" - sseEventParseDone ParseEventName = "parse_done" - sseEventParseFailed ParseEventName = "parse_failed" + SSEEventParseQueued ParseEventName = "parse_queued" + SSEEventParseUpdate ParseEventName = "parse_update" + SSEEventParseDone ParseEventName = "parse_done" + SSEEventParseFailed ParseEventName = "parse_failed" ) func (c *Client) handleSSEResponse(req *http.Request, onUpdate ParseResultUpdateFunc) (*ParseResult, error) { @@ -151,19 +151,19 @@ func (c *Client) handleSSEResponse(req *http.Request, onUpdate ParseResultUpdate } switch ev.Name() { - case string(sseEventParseQueued), string(sseEventParseUpdate): + case string(SSEEventParseQueued), string(SSEEventParseUpdate): if onUpdate != nil { onUpdate(ParseEventName(ev.Name()), &result) } continue - case string(sseEventParseDone): + case string(SSEEventParseDone): if onUpdate != nil { onUpdate(ParseEventName(ev.Name()), &result) } return &result, nil - case string(sseEventParseFailed): + case string(SSEEventParseFailed): if onUpdate != nil { onUpdate(ParseEventName(ev.Name()), &result) } diff --git a/parse_get_test.go b/parse_get_test.go index df61b57..010d032 100644 --- a/parse_get_test.go +++ b/parse_get_test.go @@ -71,7 +71,7 @@ func TestGetParseResultSSE(t *testing.T) { // Validate parse results. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) } diff --git a/parse_list.go b/parse_list.go index df67f94..e903058 100644 --- a/parse_list.go +++ b/parse_list.go @@ -25,14 +25,14 @@ import ( ) // IterParseJobs iterates over all parse jobs in the project. -func (c *Client) IterParseJobs(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] { +func (c *Client) IterParseJobs(ctx context.Context, batchSize int) iter.Seq2[ParseResult, error] { return func(yield func(ParseResult, error) bool) { cursor := "" for { listResp, err := c.ListParseJobs(ctx, &ListParseJobsRequest{ Cursor: cursor, - Limit: limit, - Direction: direction, + Limit: batchSize, + Direction: PaginationDirectionNext, }) if err != nil { yield(ParseResult{}, err) diff --git a/parse_parse_test.go b/parse_parse_test.go index 5cc3e5e..e6b89f4 100644 --- a/parse_parse_test.go +++ b/parse_parse_test.go @@ -76,7 +76,7 @@ func TestParseDocumentRemote(t *testing.T) { // Validate parse results. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) } @@ -212,7 +212,7 @@ func TestParseDocumentStructuredExtraction(t *testing.T) { // Validate parse results. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) } diff --git a/util_test.go b/util_test.go index 4f59adf..95df4eb 100644 --- a/util_test.go +++ b/util_test.go @@ -28,7 +28,7 @@ func testCleanupFileAndParseJob(t *testing.T, c *Client, fileId string, parseId // Check if file is deleted. files := []string{} - for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) { + for f, err := range c.IterFiles(t.Context(), 1) { if err != nil { t.Fatalf("failed to list files: %v", err) } @@ -48,7 +48,7 @@ func testCleanupFileAndParseJob(t *testing.T, c *Client, fileId string, parseId // Check if parse job is deleted. jobs := []string{} - for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) { + for j, err := range c.IterParseJobs(t.Context(), 1) { if err != nil { t.Fatalf("failed to list parse jobs: %v", err) }