import tempfile import unittest from pathlib import Path from unittest.mock import Mock, patch import pymupdf from PIL import Image from copienator import atomic_write_json, configuration, read_json from copienator.annotation_layout import ( answer_parts, copy_answer_pdf, fingerprint, layout_for, legacy_layout, move_answer_pdf, save_parts, added_margin, ) from copienator.commands import annotating, annotating_by_label, splitting_int from copienator.commands.annotating_with_checks import CheckboxRenderer from copienator.commands import reading_grouped_annotations as reader from copienator.left_margin import analyse_margins from copienator.pdf_cut import split_pdf class AnnotationLayoutTests(unittest.TestCase): def setUp(self): self.temp = tempfile.TemporaryDirectory() self.addCleanup(self.temp.cleanup) self.root = Path(self.temp.name) self.pdf = self.root / 'Copies' / 'Copie01.pdf' self.pdf.parent.mkdir() with pymupdf.open() as document: for _ in range(2): page = document.new_page(width=600, height=800) page.insert_text((100, 100), 'Answer') document.save(self.pdf) def test_all_pages_measured_after_early_no_margin_run(self): images = [Image.new('RGB', (600, 100)) for _ in range(6)] values = [10, 10, 40, None, 19.9, 20] estimates = {id(image): value for image, value in zip(images, values)} def estimate(image, dpi): edge = estimates[id(image)] return {'edge_mm': edge, 'bands': 0 if edge is None else 8, 'band_starts_mm': [] if edge is None else [edge] * 8} with patch('copienator.left_margin.estimate_page', side_effect=estimate): result = analyse_margins(images, 200, 600) self.assertEqual([p['status'] for p in result['pages']], ['low', 'low', 'normal', 'unknown', 'low', 'low']) self.assertAlmostEqual(result['pages'][3]['effective_edge_mm'], 19.98) def test_split_records_pages_and_rejects_stale_measurements(self): schema_path = self.root / 'Cutleft' / 'Copie01_schema.json' atomic_write_json(schema_path, { 'source_sha256': fingerprint(self.pdf), 'auto_margin': {'pages': [{'page': 1, 'edge_mm': 40}, {'page': 2, 'edge_mm': 10}]}, }) staging = self.root / 'split' staging.mkdir() splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging) answer = staging / 'Ex 1.pdf' self.assertEqual([p['page'] for p in answer_parts(answer)], [1, 2]) self.assertEqual(layout_for(answer)['page_margins'], [300, 450]) atomic_write_json(schema_path, {'source_sha256': 'stale', 'auto_margin': { 'pages': [{'page': 1, 'edge_mm': 10}]}}) splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging) self.assertEqual(layout_for(answer)['margin_left'], 300) def test_configurable_margin_unknown_fallback_and_stale_sidecar(self): self.assertEqual(layout_for(self.pdf)['margin_left'], 300) save_parts(self.pdf, [{'margin_status': 'normal', 'margin_mm': 40}, {'margin_status': 'low', 'margin_mm': 20}]) with patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 500), (20, 450), (35, 300)]): self.assertEqual(layout_for(self.pdf)['margin_left'], 450) self.pdf.write_bytes(self.pdf.read_bytes() + b'\n') self.assertEqual(layout_for(self.pdf)['margin_left'], 300) def test_horizontal_answer_crop_can_remove_an_original_margin(self): atomic_write_json(self.root / 'Cutleft' / 'Copie01_schema.json', { 'source_sha256': fingerprint(self.pdf), 'auto_margin': {'pages': [{'page': 1, 'edge_mm': 40}, {'page': 2, 'edge_mm': 40}]}, }) staging = self.root / 'horizontal' staging.mkdir() # 20% of a 600-point page removes 42.3 mm from the left. splitting_int._render_split_outputs(self.pdf, [('|Ex 1', 0, 0, 10, 100, 110)], staging) answer = staging / 'Ex 1.pdf' self.assertEqual(layout_for(answer)['margin_left'], 450) self.assertEqual(answer_parts(answer)[0]['margin_mm'], 0) def test_copy_move_and_vertical_cut_preserve_page_classification(self): save_parts(self.pdf, [{'margin_status': 'normal', 'margin_mm': 40}, {'margin_status': 'low', 'margin_mm': 20}]) copied = self.root / 'copied.pdf' moved = self.root / 'moved.pdf' copy_answer_pdf(self.pdf, copied) move_answer_pdf(copied, moved) self.assertEqual(layout_for(moved)['margin_left'], 400) first, second = self.root / 'first.pdf', self.root / 'second.pdf' split_pdf(moved, 50, first, second) self.assertEqual(layout_for(first)['margin_left'], 300) self.assertEqual(layout_for(second)['margin_left'], 400) def render(self, layout): callback = Mock() result = {'score': 2, 'feedback': [{'text': 'Feedback', 'box_2d': [10, 50, 30, 100]}]} with patch.object(annotating, 'render_score_text', return_value=Image.new('RGBA', (200, 40))): image, header = annotating.compose_label_image( Image.new('RGB', (1654, 100), 'blue'), 'Ex 1', result, 0, render_fn=lambda *args, **kwargs: Image.new('RGBA', (200, 10)), draw_callback=callback, layout=layout) box = next(call.args[2]['box'] for call in callback.call_args_list if call.args[0] == 'local_rect') return image, header, box def test_interpolation_and_clamping(self): for mm, pixels in [(None, 300), (0, 450), (10, 450), (15, 425), (20, 400), (27.5, 350), (35, 300), (80, 300)]: with self.subTest(mm=mm): self.assertEqual(added_margin(mm), pixels) def test_unknown_measurement_inherits_average_during_answer_extraction(self): atomic_write_json(self.root / 'Cutleft' / 'Copie01_schema.json', { 'source_sha256': fingerprint(self.pdf), 'auto_margin': {'pages': [{'page': 1, 'edge_mm': 15}, {'page': 2, 'edge_mm': None}]}, }) staging = self.root / 'inferred' staging.mkdir() splitting_int._render_split_outputs(self.pdf, [('Ex 1', 0, 0, 10, 0, 10)], staging) answer = staging / 'Ex 1.pdf' self.assertEqual(answer_parts(answer)[1]['margin_status'], 'unknown') self.assertEqual(answer_parts(answer)[1]['margin_mm'], 15) self.assertEqual(layout_for(answer)['page_margins'], [425, 425]) def test_per_page_offsets_narrow_pages_and_saved_layout(self): save_parts(self.pdf, [{'margin_mm': 10}, {'margin_mm': 35}]) pages = [Image.new('RGB', (1000, 100), 'blue'), Image.new('RGB', (800, 100), 'green')] with patch.object(annotating, 'convert_from_path', return_value=pages): base, _, _ = annotating.make_base_image(self.pdf) result = {'score': 2, 'feedback': [{'text': 'Across pages', 'box_2d': [80, 150, 120, 200]}]} callback = Mock() with patch.object(annotating, 'render_score_text', return_value=Image.new('RGB', (200, 40))): def render(**kwargs): return annotating.compose_label_image( base, 'Ex 1', result, 0, render_fn=lambda *a, **k: Image.new('RGBA', (10, 10)), **kwargs) image, header = render(source_pdf=self.pdf, draw_callback=callback) boxes = [call.args[2]['box'] for call in callback.call_args_list if call.args[0] == 'local_rect'] self.assertEqual(boxes, [[600, 120, 650, 140], [350, 140, 400, 160]]) for x, y, color in [(450, header + 20, (0, 0, 255)), (300, header + 150, (0, 128, 0))]: self.assertEqual(image.getpixel((x, y)), color) self.assertEqual(image.getpixel((x - 1, y)), (255, 255, 255)) with patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 900)]): restored, _ = render(layout=image.info['annotation_layout']) self.assertEqual(restored.tobytes(), image.tobytes()) def test_wider_layout_translates_copy_and_feedback_without_scaling(self): normal, height, box = self.render(legacy_layout()) low, low_height, low_box = self.render({**legacy_layout(), 'margin_left': 400}) self.assertEqual(low.width, normal.width + 100) self.assertEqual(height, low_height) self.assertEqual(low_box, [box[0] + 100, box[1], box[2] + 100, box[3]]) self.assertEqual(low.getpixel((400, height + 80)), (0, 0, 255)) self.assertEqual(low.getpixel((399, height + 80)), (255, 255, 255)) def test_later_rectangle_cannot_draw_inside_comment_checkbox(self): renderer = CheckboxRenderer('Ex 1') result = {'score': 2, 'feedback': [ {'text': 'First comment', 'box_2d': [10, 50, 50, 200]}, {'text': 'Later rectangle', 'box_2d': [10, 300, 80, 400]}, ]} with patch.object(annotating, 'render_score_text', return_value=Image.new('RGB', (800, 80), 'white')): image, _ = annotating.compose_label_image( Image.new('RGB', (1654, 100), 'white'), 'Ex 1', result, 0, render_fn=lambda *args, **kwargs: Image.new('RGBA', (620, 50), (255, 200, 200, 180)), draw_callback=renderer.callback, layout=legacy_layout()) # The later rectangle's left edge crosses the first comment's checkbox. checkbox = next(item for item in renderer.checkboxes if item['type'] == 'del_local' and item['index'] == 0) x1, y1, x2, y2 = checkbox['final_box'] interior = image.crop((x1 + 5, y1 + 5, x2 - 5, y2 - 5)) self.assertEqual(interior.getextrema(), ((255, 255),) * 3) # The rectangle still exists outside the checkbox. self.assertEqual(image.getpixel((650, 90)), (255, 0, 0)) def test_mixed_bgnot_widths_and_saved_geometry_survive_config_change(self): normal, height, _ = self.render(legacy_layout()) low, _, _ = self.render({**legacy_layout(), 'margin_left': 400}) annotating_by_label.save_batch([ ('01', 'Ex 1', normal, height, []), ('02', 'Ex 1', low, height, []), ], 'Ex 1', 1, self.root) output = self.root / 'Ex 1 G1' meta = read_json(output / 'bnote.json') self.assertEqual(meta['width'], low.width) with Image.open(output / 'Reference.jpg') as reference: self.assertGreater(min(reference.getpixel((low.width - 10, 80))), 245) with patch.object(reader, 'detect_checks_and_notes', return_value=([], Image.new('RGBA', (low.width, low.height * 2)))), patch.object(configuration, 'ANNOTATION_MARGIN_POINTS', [(10, 500)]): _, notes = reader._scan_annotation_directory(output) self.assertEqual(notes['02']['Ex 1']['layout']['margin_left'], 400) rerender, _, _ = self.render(notes['02']['Ex 1']['layout']) self.assertEqual(rerender.width, low.width) with patch.object(reader, 'detect_checks_and_notes', return_value=([], None)): _, notes = reader._scan_annotation_directory(output) self.assertEqual(notes['02']['Ex 1']['layout']['margin_left'], 400) self.assertIsNone(notes['02']['Ex 1']['img']) if __name__ == '__main__': unittest.main()