diff --git a/tests/parse/test_parse_vrplib.py b/tests/parse/test_parse_vrplib.py index f7473666..0c040c13 100644 --- a/tests/parse/test_parse_vrplib.py +++ b/tests/parse/test_parse_vrplib.py @@ -114,8 +114,8 @@ def test_parse_specification(line, key, value): ["demand", np.array([1.1, 2.2, 3.3])], ), ( - ["DEPOT_SECTION", "1", "-1"], - ["depot", np.array([0])], + ["DEPOT_SECTION", "1", "2", "3", "-1"], + ["depot", np.array([0, 1, 2])], ), ( # No end token in DEPOT_SECTION diff --git a/vrplib/parse/parse_vrplib.py b/vrplib/parse/parse_vrplib.py index d3533529..9cfd1640 100644 --- a/vrplib/parse/parse_vrplib.py +++ b/vrplib/parse/parse_vrplib.py @@ -111,20 +111,21 @@ def parse_section( Parses the data section lines. """ name = lines[0].strip().removesuffix("_SECTION").lower() - values = [[infer_type(n) for n in line.split()] for line in lines[1:]] + rows = [[infer_type(n) for n in line.split()] for line in lines[1:]] if name == "edge_weight": # Parse edge weights separately as it involves extra processing. - data = parse_distances(values, **instance) # type: ignore + data = parse_distances(rows, **instance) # type: ignore elif name == "depot": # Remove -1 end token and renormalize depots to start at zero. - data = np.array(values[0]) - 1 - elif any(len(row) != len(values[0]) for row in values): + data = np.array(rows) + data = data[data != -1] - 1 + elif any(len(row) != len(rows[0]) for row in rows): # This is a ragged array, so we keep it as a nested list, but we # remove the indices column. - data = [row[1:] for row in values] + data = [row[1:] for row in rows] else: - data = np.array([row[1:] for row in values]) + data = np.array([row[1:] for row in rows]) if data.ndim > 1 and data.shape[-1] == 1: # Squeeze data lines that contain only one column.