Skip to content

Commit 5395f6c

Browse files
OUDONOUDON
committed
feat: SegmentOptimizer#compute raises DataTooLongError
1 parent 8b08545 commit 5395f6c

File tree

3 files changed

+46
-31
lines changed

3 files changed

+46
-31
lines changed

‎src/rmqrcode/rmqrcode.py

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -82,23 +82,24 @@ def fit(data, ecc=ErrorCorrectionLevel.M, fit_strategy=FitStrategy.BALANCED):
8282
logger.debug("Select rMQR Code version")
8383
for version_name, qr_version in DataCapacities.items():
8484
optimizer = qr_segments.SegmentOptimizer()
85-
optimized_segments = optimizer.compute(data, version_name)
86-
data_length = qr_segments.compute_length(optimized_segments, version_name)
87-
88-
if data_length <= qr_version["number_of_data_bits"][ecc]:
89-
width, height = qr_version["width"], qr_version["height"]
90-
if width not in determined_width and height not in determined_height:
91-
determined_width.add(width)
92-
determined_height.add(height)
93-
ok_versions.append(
94-
{
95-
"version": version_name,
96-
"width": width,
97-
"height": height,
98-
"segments": optimized_segments,
99-
}
100-
)
101-
logger.debug(f"ok: {version_name}")
85+
try:
86+
optimized_segments = optimizer.compute(data, version_name, ecc)
87+
except DataTooLongError:
88+
continue
89+
90+
width, height = qr_version["width"], qr_version["height"]
91+
if width not in determined_width and height not in determined_height:
92+
determined_width.add(width)
93+
determined_height.add(height)
94+
ok_versions.append(
95+
{
96+
"version": version_name,
97+
"width": width,
98+
"height": height,
99+
"segments": optimized_segments,
100+
}
101+
)
102+
logger.debug(f"ok: {version_name}")
102103

103104
if len(ok_versions) == 0:
104105
raise DataTooLongError("The data is too long.")
@@ -128,7 +129,7 @@ def sort_key(x):
128129

129130
def _optimized_segments(self, data):
130131
optimizer = qr_segments.SegmentOptimizer()
131-
return optimizer.compute(data, self.version_name())
132+
return optimizer.compute(data, self.version_name(), self._error_correction_level)
132133

133134
def __init__(self, version, ecc, with_quiet_zone=True, logger=None):
134135
self._logger = logger or rMQR._init_logger()

‎src/rmqrcode/segments.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from . import encoder
22
from .errors import DataTooLongError
33
from .format.rmqr_versions import rMQRVersions
4+
from .format.data_capacities import DataCapacities
45

56
encoders = [
67
encoder.NumericEncoder,
@@ -47,12 +48,13 @@ def __init__(self):
4748
self.dp = [[[self.INF for n in range(3)] for mode in range(4)] for length in range(self.MAX_CHARACTER + 1)]
4849
self.parents = [[[-1 for n in range(3)] for mode in range(4)] for length in range(self.MAX_CHARACTER + 1)]
4950

50-
def compute(self, data, version):
51+
def compute(self, data, version, ecc):
5152
"""Computes the optimize segmentation for the given data.
5253
5354
Args:
5455
data (str): The data to encode.
5556
version (str): The version name.
57+
ecc (rmqrcode.ErrorCorrectionLevel): The error correction level.
5658
5759
Returns:
5860
list: The list of segments.
@@ -66,8 +68,11 @@ def compute(self, data, version):
6668

6769
self.qr_version = rMQRVersions[version]
6870
self._compute_costs(data)
69-
best_index = self._find_best(data)
70-
path = self._reconstruct_path(best_index)
71+
best = self._find_best(data)
72+
if best["cost"] > DataCapacities[version]["number_of_data_bits"][ecc]:
73+
raise DataTooLongError
74+
75+
path = self._reconstruct_path(best["index"])
7176
segments = self._compute_segments(path, data)
7277
return segments
7378

@@ -139,7 +144,8 @@ def _find_best(self, data):
139144
data (str): The data to encode.
140145
141146
Returns:
142-
tuple: The best index as tuple (n, mode, unfilled_length).
147+
dict: The dict object includes "cost" and "index". The "cost" is the value of minimum cost.
148+
The "index" is the index of the dp table as a tuple (n, mode, unfilled_length).
143149
144150
"""
145151
best = self.INF
@@ -149,7 +155,10 @@ def _find_best(self, data):
149155
if self.dp[len(data)][mode][unfilled_length] < best:
150156
best = self.dp[len(data)][mode][unfilled_length]
151157
best_index = (len(data), mode, unfilled_length)
152-
return best_index
158+
return {
159+
"cost": best,
160+
"index": best_index
161+
}
153162

154163
def _reconstruct_path(self, best_index):
155164
"""Reconstructs the path.

‎tests/segments_test.py

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,54 +1,59 @@
11
from rmqrcode.segments import SegmentOptimizer, compute_length
2-
from rmqrcode import encoder
2+
from rmqrcode import encoder, ErrorCorrectionLevel, DataTooLongError
33
import pytest
44

55

66
class TestSegments:
77
def test_can_optimize_segments_numeric_and_byte(self):
88
optimizer = SegmentOptimizer()
9-
segments = optimizer.compute("123Abc", "R7x43")
9+
segments = optimizer.compute("123Abc", "R7x43", ErrorCorrectionLevel.M)
1010
assert segments == [
1111
{"data": "123", "encoder_class": encoder.NumericEncoder},
1212
{"data": "Abc", "encoder_class": encoder.ByteEncoder},
1313
]
1414

1515
def test_can_optimize_segments_alphanumeric_and_kanji(self):
1616
optimizer = SegmentOptimizer()
17-
segments = optimizer.compute("17:30集合", "R7x59")
17+
segments = optimizer.compute("17:30集合", "R7x59", ErrorCorrectionLevel.M)
1818
assert segments == [
1919
{"data": "17:30", "encoder_class": encoder.AlphanumericEncoder},
2020
{"data": "集合", "encoder_class": encoder.KanjiEncoder},
2121
]
2222

2323
def test_can_optimize_segments_numeric_only(self):
2424
optimizer = SegmentOptimizer()
25-
segments = optimizer.compute("123456", "R7x59")
25+
segments = optimizer.compute("123456", "R7x59", ErrorCorrectionLevel.M)
2626
assert segments == [
2727
{"data": "123456", "encoder_class": encoder.NumericEncoder},
2828
]
2929

3030
def test_can_optimize_segments_alphanumeric_only(self):
3131
optimizer = SegmentOptimizer()
32-
segments = optimizer.compute("HTTPS://", "R7x59")
32+
segments = optimizer.compute("HTTPS://", "R7x59", ErrorCorrectionLevel.M)
3333
assert segments == [
3434
{"data": "HTTPS://", "encoder_class": encoder.AlphanumericEncoder},
3535
]
3636

3737
def test_can_optimize_segments_byte_only(self):
3838
optimizer = SegmentOptimizer()
39-
segments = optimizer.compute("1+zY!a:K", "R7x59")
39+
segments = optimizer.compute("1+zY!a:K", "R7x59", ErrorCorrectionLevel.M)
4040
assert segments == [
4141
{"data": "1+zY!a:K", "encoder_class": encoder.ByteEncoder},
4242
]
4343

4444
def test_can_optimize_segments_kanji_only(self):
4545
optimizer = SegmentOptimizer()
46-
segments = optimizer.compute("漢字", "R7x59")
46+
segments = optimizer.compute("漢字", "R7x59", ErrorCorrectionLevel.M)
4747
assert segments == [
4848
{"data": "漢字", "encoder_class": encoder.KanjiEncoder},
4949
]
5050

51+
def test_optimize_segments_raises_data_too_long_error(self):
52+
optimizer = SegmentOptimizer()
53+
with pytest.raises(DataTooLongError) as e:
54+
segments = optimizer.compute("a" * 12, "R7x59", ErrorCorrectionLevel.M)
55+
5156
def test_compute_length(self):
5257
optimizer = SegmentOptimizer()
53-
segments = optimizer.compute("123Abc", "R7x43")
58+
segments = optimizer.compute("123Abc", "R7x43", ErrorCorrectionLevel.M)
5459
assert compute_length(segments, "R7x43") is 47

0 commit comments

Comments
 (0)