Merged in feature/mergecell (pull request #142)

Merged Cell

* mergecellv1

* buildimg

* merging

* addfile

* justmergedcell

internal/document/text/childElements.go

@@ -51,7 +51,7 @@ func (s *Service) getBlockMap(blocksList []types.Block) map[uuid.UUID]types.Bloc
 		}
 		id, err := uuid.Parse(*block.Id)
 		if err != nil {
-			slog.Warn("error parsing block id", "error", err, "id", *block.Id)
+			slog.Warn("error parsing block id for block map", "error", err, "id", *block.Id)
 			continue
 		}
 

internal/document/text/documentText_test.go

@@ -67,7 +67,7 @@ func TestGetDocumentText(t *testing.T) {
 		require.EqualError(t, err, "no textract response")
 	})
 
-	t.Run("sampleOne - cnc_01-06", func(t *testing.T) {
+	t.Run("cnc_01-06", func(t *testing.T) {
 		t.Parallel()
 		cfg := DocTextConfig{}
 		svc := Service{
@@ -86,7 +86,7 @@ func TestGetDocumentText(t *testing.T) {
 
 		assertReaders(t, txtFile, text)
 	})
-	t.Run("sampleOne - hn_0109", func(t *testing.T) {
+	t.Run("hn_0109", func(t *testing.T) {
 		t.Parallel()
 		cfg := DocTextConfig{}
 		svc := Service{
@@ -105,7 +105,7 @@ func TestGetDocumentText(t *testing.T) {
 
 		assertReaders(t, txtFile, text)
 	})
-	t.Run("sampleOne - hn_23-70", func(t *testing.T) {
+	t.Run("hn_23-70", func(t *testing.T) {
 		t.Parallel()
 		cfg := DocTextConfig{}
 		svc := Service{
@@ -124,7 +124,7 @@ func TestGetDocumentText(t *testing.T) {
 
 		assertReaders(t, txtFile, text)
 	})
-	t.Run("sampleOne - chc_1", func(t *testing.T) {
+	t.Run("chc_1", func(t *testing.T) {
 		t.Parallel()
 		cfg := DocTextConfig{}
 		svc := Service{
@@ -141,6 +141,25 @@ func TestGetDocumentText(t *testing.T) {
 		text, err := svc.getDocumentText(ctx, pdf)
 		require.NoError(t, err)
 
+		assertReaders(t, txtFile, text)
+	})
+	t.Run("merged_cell_table", func(t *testing.T) {
+		t.Parallel()
+		cfg := DocTextConfig{}
+		svc := Service{
+			cfg: &cfg,
+		}
+		mockTextract := textractmock.NewMockTextractClient(t)
+		cfg.TextractClient = mockTextract
+
+		textractBaseOut, pdf, txtFile, close := getFiles(t, "merged_cell_table")
+		defer close()
+
+		textExpectations(t, ctx, mockTextract, pdf, textractBaseOut)
+
+		text, err := svc.getDocumentText(ctx, pdf)
+		require.NoError(t, err)
+
 		assertReaders(t, txtFile, text)
 	})
 }

internal/document/text/pageText.go

@@ -32,15 +32,15 @@ func (s *Service) getPageText(ctx context.Context, document documenttypes.File,
 
 	blockMap := s.getBlockMap(res.Blocks)
 
-	lines, tables, tableBlockMap, err := s.getPageElements(ctx, document, index, *pageBlock, blockMap)
+	lines, specialFeatures, err := s.getPageElements(ctx, document, index, *pageBlock, blockMap)
 	if err != nil {
 		return "", err
 	}
 
-	return s.buildPageText(ctx, lines, blockMap, index, tables, tableBlockMap)
+	return s.buildPageText(ctx, lines, blockMap, index, specialFeatures)
 }
 
-func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, pageBlock types.Block, blockMap map[uuid.UUID]types.Block) ([]uuid.UUID, []uuid.UUID, map[uuid.UUID]types.Block, error) {
+func (s *Service) getPageElements(ctx context.Context, document documenttypes.File, index int, pageBlock types.Block, blockMap map[uuid.UUID]types.Block) ([]uuid.UUID, specialisedFeatures, error) {
 	directChildren := s.getDirectChildren(pageBlock, blockMap)
 
 	hasTable := false
@@ -50,18 +50,30 @@ func (s *Service) getPageElements(ctx context.Context, document documenttypes.Fi
 		}
 	}
 
-	var tables []uuid.UUID
-	var tableBlockMap map[uuid.UUID]types.Block
+	features := []types.FeatureType{}
 	if hasTable {
-		ids, blockMap, err := s.getTables(ctx, document, index)
-		if err != nil {
-			return nil, nil, nil, err
-		}
-		tables = ids
-		tableBlockMap = blockMap
+		features = append(features, types.FeatureTypeTables)
+		// features = append(features, types.FeatureTypeForms)
 	}
 
-	return directChildren, tables, tableBlockMap, nil
+	// Detecting Checkboxes -
+	// if has tables
+	// For a page with no tables, we look for 3 out of 4 of the following conditions
+	//     has_exhibit→ 'exhibit' is on the page
+	//     has_benefit→ one of ['benefit plan', 'benefit plan(s)', 'applicable benefit', 'benefits'] is on the page
+	//     has_plans→ one of ['plan(s):', 'plan(s)/program(s):', 'plans:', 'plan/program:', 'plan(s)', 'plans'] is on the page
+	//     has_services → one of ['services:', 'service:', 'service(s):', 'services', 'service(s)'] is on the page
+
+	if len(features) == 0 {
+		return directChildren, specialisedFeatures{}, nil
+	}
+
+	featuresOut, err := s.getSpecialisedFeatures(ctx, document, index, features)
+	if err != nil {
+		return nil, specialisedFeatures{}, err
+	}
+
+	return directChildren, featuresOut, nil
 }
 
 type buildParams struct {
@@ -69,7 +81,7 @@ type buildParams struct {
 	currentTable    int
 }
 
-func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap map[uuid.UUID]types.Block, index int, tables []uuid.UUID, tableBlockMap map[uuid.UUID]types.Block) (string, error) {
+func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap map[uuid.UUID]types.Block, index int, specialFeatures specialisedFeatures) (string, error) {
 	var builder strings.Builder
 	params := buildParams{
 		numOfSignatures: 0,
@@ -77,15 +89,15 @@ func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap
 	}
 
 	for _, id := range lines {
-		err := s.addComponentText(ctx, blockMap[id], blockMap, tables, tableBlockMap, &params, &builder)
+		err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, &params, &builder)
 		if err != nil {
 			return "", err
 		}
 	}
 
-	for params.currentTable < len(tables) {
-		id := tables[params.currentTable]
-		err := s.addComponentText(ctx, tableBlockMap[id], blockMap, tables, tableBlockMap, &params, &builder)
+	for params.currentTable < len(specialFeatures.Tables) {
+		id := specialFeatures.Tables[params.currentTable]
+		err := s.addComponentText(ctx, specialFeatures.BlockMap[id], blockMap, specialFeatures, &params, &builder)
 		if err != nil {
 			return "", err
 		}
@@ -98,11 +110,11 @@ func (s *Service) buildPageText(ctx context.Context, lines []uuid.UUID, blockMap
 	return builder.String(), nil
 }
 
-func (s *Service) addComponentText(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, tables []uuid.UUID, tableBlockMap map[uuid.UUID]types.Block, params *buildParams, builder *strings.Builder) error {
+func (s *Service) addComponentText(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, specialFeatures specialisedFeatures, params *buildParams, builder *strings.Builder) error {
 	switch block.BlockType {
 	case types.BlockTypeLayoutText, types.BlockTypeLayoutTitle, types.BlockTypeLayoutHeader, types.BlockTypeLayoutFooter, types.BlockTypeLayoutSectionHeader, types.BlockTypeLayoutPageNumber, types.BlockTypeLayoutList, types.BlockTypeLayoutFigure, types.BlockTypeLayoutKeyValue:
 		for _, id := range s.getChildIds(block) {
-			err := s.addComponentText(ctx, blockMap[id], blockMap, tables, tableBlockMap, params, builder)
+			err := s.addComponentText(ctx, blockMap[id], blockMap, specialFeatures, params, builder)
 			if err != nil {
 				return err
 			}
@@ -122,9 +134,9 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 	case types.BlockTypeSignature:
 		params.numOfSignatures++
 	case types.BlockTypeLayoutTable, types.BlockTypeTable:
-		if block.BlockType == types.BlockTypeLayoutTable && len(tables) <= params.currentTable {
+		if block.BlockType == types.BlockTypeLayoutTable && len(specialFeatures.Tables) <= params.currentTable {
 			for _, child := range s.getChildIds(block) {
-				err := s.addComponentText(ctx, blockMap[child], blockMap, tables, tableBlockMap, params, builder)
+				err := s.addComponentText(ctx, blockMap[child], blockMap, specialFeatures, params, builder)
 				if err != nil {
 					return err
 				}
@@ -132,10 +144,10 @@ func (s *Service) addComponentText(ctx context.Context, block types.Block, block
 			params.currentTable++
 			break
 		}
-		tableId := tables[params.currentTable]
+		tableId := specialFeatures.Tables[params.currentTable]
 		params.currentTable++
-		tableBlock := tableBlockMap[tableId]
-		err := s.addTable(ctx, tableBlock, tableBlockMap, builder)
+		tableBlock := specialFeatures.BlockMap[tableId]
+		err := s.addTable(ctx, tableBlock, specialFeatures.BlockMap, builder)
 		if err != nil {
 			return err
 		}

internal/document/text/pageText_test.go

@@ -77,7 +77,7 @@ func TestGetBlockMap(t *testing.T) {
 		}, outBlocks)
 	})
 	t.Run("invalid block id", func(t *testing.T) {
-		blockIdStr := "blockId.String()"
+		blockIdStr := "invalid_block_id"
 		blocks := []types.Block{
 			{
 				Id: &blockIdStr,

internal/document/text/tables.go

@@ -32,9 +32,9 @@ func (s *Service) addTable(ctx context.Context, block types.Block, blockMap map[
 		s.addComponentTable(ctx, blockMap[uid], blockMap, &params)
 	}
 
-	tableArr := make([][]string, params.maxRow)
+	tableArr := make([][]string, params.maxRow+1)
 	for i := range tableArr {
-		tableArr[i] = make([]string, params.maxColumn)
+		tableArr[i] = make([]string, params.maxColumn+1)
 	}
 
 	for _, cell := range params.table {
@@ -62,53 +62,91 @@ type tableParams struct {
 func (s *Service) addComponentTable(ctx context.Context, block types.Block, blockMap map[uuid.UUID]types.Block, params *tableParams) {
 	switch block.BlockType {
 	case types.BlockTypeMergedCell:
-		for _, id := range s.getChildIds(block) {
-			s.addComponentTable(ctx, blockMap[id], blockMap, params)
+		content := ""
+		for _, uid := range s.getChildIds(block) {
+			if blockMap[uid].BlockType != types.BlockTypeCell {
+				slog.Warn("not a cell", "block_type", blockMap[uid].BlockType)
+				continue
+			}
+			content = s.getCellContent(blockMap[uid], blockMap)
+			if content != "" {
+				break
+			}
+		}
+		startColumnIndex := int(*block.ColumnIndex - 1)
+		startRowIndex := int(*block.RowIndex - 1)
+		endColumnIndex := startColumnIndex + int(*block.ColumnSpan) - 1
+		endRowIndex := startRowIndex + int(*block.RowSpan) - 1
+		for row := startRowIndex; row <= endRowIndex; row++ {
+			for column := startColumnIndex; column <= endColumnIndex; column++ {
+				s.addCellToTable(
+					row,
+					column,
+					content,
+					params,
+				)
+			}
 		}
 	case types.BlockTypeCell:
-		var str strings.Builder
-		for _, uid := range s.getChildIds(block) {
-			if blockMap[uid].BlockType == types.BlockTypeWord {
-				if str.String() != "" {
-					str.WriteRune(' ')
-				}
-				str.WriteString(*blockMap[uid].Text)
-			}
-
-			if params.maxRow < int(*block.RowIndex) {
-				params.maxRow = int(*block.RowIndex)
-			}
-			if params.maxColumn < int(*block.ColumnIndex) {
-				params.maxColumn = int(*block.ColumnIndex)
-			}
-
-			params.table = append(params.table, cell{
-				column: int(*block.ColumnIndex) - 1,
-				row:    int(*block.RowIndex) - 1,
-				value:  str.String(),
-			})
-		}
+		content := s.getCellContent(block, blockMap)
+		s.addCellToTable(
+			int(*block.RowIndex)-1,
+			int(*block.ColumnIndex)-1,
+			content,
+			params,
+		)
 	default:
 		slog.Info("unhandled block type in table", "type", block.BlockType)
 	}
 }
 
-func (s *Service) getTables(ctx context.Context, document documenttypes.File, index int) ([]uuid.UUID, map[uuid.UUID]types.Block, error) {
+func (s *Service) getCellContent(block types.Block, blockMap map[uuid.UUID]types.Block) string {
+	var str strings.Builder
+	for _, uid := range s.getChildIds(block) {
+		if blockMap[uid].BlockType == types.BlockTypeWord {
+			if str.String() != "" {
+				str.WriteRune(' ')
+			}
+			str.WriteString(*blockMap[uid].Text)
+		}
+	}
+
+	return str.String()
+}
+
+func (s *Service) addCellToTable(rowIndex int, columnIndex int, content string, params *tableParams) {
+	if params.maxRow < rowIndex {
+		params.maxRow = rowIndex
+	}
+	if params.maxColumn < columnIndex {
+		params.maxColumn = columnIndex
+	}
+	params.table = append(params.table, cell{
+		column: columnIndex,
+		row:    rowIndex,
+		value:  content,
+	})
+}
+
+type specialisedFeatures struct {
+	BlockMap map[uuid.UUID]types.Block
+	Tables   []uuid.UUID
+}
+
+func (s *Service) getSpecialisedFeatures(ctx context.Context, document documenttypes.File, index int, features []types.FeatureType) (specialisedFeatures, error) {
 	pageBytes, err := document.GetPageAsPNG(ctx, index)
 	if err != nil {
-		return nil, nil, err
+		return specialisedFeatures{}, err
 	}
 
 	res, err := s.cfg.GetTextractClient().AnalyzeDocument(ctx, &textract.AnalyzeDocumentInput{
 		Document: &types.Document{
 			Bytes: pageBytes,
 		},
-		FeatureTypes: []types.FeatureType{
-			types.FeatureTypeTables,
-		},
+		FeatureTypes: features,
 	})
 	if err != nil {
-		return nil, nil, err
+		return specialisedFeatures{}, err
 	}
 
 	tables := []uuid.UUID{}
@@ -116,7 +154,7 @@ func (s *Service) getTables(ctx context.Context, document documenttypes.File, in
 		if block.BlockType == types.BlockTypeTable {
 			id, err := uuid.Parse(*block.Id)
 			if err != nil {
-				slog.Warn("error parsing block id", "error", err, "id", *block.Id)
+				slog.Warn("error parsing block id for tables", "error", err, "id", *block.Id)
 				continue
 			}
 
@@ -125,5 +163,8 @@ func (s *Service) getTables(ctx context.Context, document documenttypes.File, in
 	}
 	blocks := s.getBlockMap(res.Blocks)
 
-	return tables, blocks, nil
+	return specialisedFeatures{
+		BlockMap: blocks,
+		Tables:   tables,
+	}, nil
 }

internal/document/types/pdf_test.go

@@ -2,10 +2,7 @@ package documenttypes
 
 import (
 	"context"
-	"io"
 	"math"
-	"os"
-	"path/filepath"
 	"strings"
 	"testing"
 
@@ -268,34 +265,6 @@ func TestIsCorrupt(t *testing.T) {
 		reason := pdf.IsCorrupt(ctx)
 		assert.Nil(t, reason)
 	})
-	t.Run("validate pdfs in assets/original", func(t *testing.T) {
-		var pdfFiles []string
-		dir := "../../../assets/original"
-
-		err := filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
-			if err != nil {
-				return err
-			}
-
-			if !info.IsDir() && strings.ToLower(filepath.Ext(path)) == ".pdf" {
-				pdfFiles = append(pdfFiles, path)
-			}
-			return nil
-		})
-		require.NoError(t, err)
-
-		for _, filename := range pdfFiles {
-			file, err := os.Open(filename)
-			require.NoError(t, err)
-			defer file.Close()
-
-			readSeeker := io.ReadSeeker(file)
-
-			pdf := NewPDF(readSeeker)
-			docErr := pdf.IsCorrupt(ctx)
-			assert.Nil(t, docErr)
-		}
-	})
 }
 
 func TestNewPDF(t *testing.T) {

internal/test/ecosystem.go

@@ -67,6 +67,12 @@ func CreateFullNetwork(t testing.TB, ctx context.Context, cfg FullDependenciesCo
 		}()
 	}
 
+	wg.Add(1)
+	go func() {
+		buildImage(t, ctx)
+		wg.Done()
+	}()
+
 	wg.Wait()
 
 	qService, err := queryapi.NewClientWithResponses(apiContainers[QueryAPIName].URI)