|
1 | | -"""Small, fail-closed helpers for the SPSS type-6 document record. |
| 1 | +"""Small, fail-closed helpers for selected raw SPSS dictionary records. |
2 | 2 |
|
3 | 3 | This is not a second SAV engine: pyspssio still owns values and the normal |
4 | 4 | dictionary. IBM I/O exposes copying document records but not their text. |
@@ -51,14 +51,155 @@ def write_document_lines(path: str | Path, lines: Iterable[str], *, encoding: st |
51 | 51 | if len(documents) > 1: |
52 | 52 | raise RawDictionaryError("SAV dictionary contains more than one document record.") |
53 | 53 | replacement = _document_record(lines, encoding=encoding, byte_order=byte_order) |
| 54 | + terminator = next(record for record in records if record.record_type == 999) |
54 | 55 | if documents: |
55 | 56 | current = documents[0] |
56 | | - target.write_bytes(data[: current.start] + replacement + data[current.end :]) |
| 57 | + updated = data[: current.start] + replacement + data[current.end :] |
| 58 | + else: |
| 59 | + extension = next((record for record in records if record.record_type == 7), None) |
| 60 | + insertion = extension.start if extension else terminator.start |
| 61 | + updated = data[:insertion] + replacement + data[insertion:] |
| 62 | + target.write_bytes(_shift_zsav_offsets( |
| 63 | + updated, |
| 64 | + original_data_start=terminator.end, |
| 65 | + delta=len(updated) - len(data), |
| 66 | + byte_order=byte_order, |
| 67 | + )) |
| 68 | + |
| 69 | + |
| 70 | +def write_compatible_names( |
| 71 | + path: str | Path, names: dict[str, str], *, encoding: str, |
| 72 | +) -> None: |
| 73 | + """Set exact legacy short names without changing long source variable names. |
| 74 | +
|
| 75 | + IBM I/O has no compatible-name setter. The standard long-name extension |
| 76 | + and the fixed 8-byte variable records contain the complete mapping, so this |
| 77 | + narrowly rewrites only those two dictionary locations. Every requested |
| 78 | + source name must have a long-name record; otherwise the operation fails |
| 79 | + instead of silently deriving another name. |
| 80 | + """ |
| 81 | + if not names: |
57 | 82 | return |
58 | | - extension = next((record for record in records if record.record_type == 7), None) |
| 83 | + target = Path(path) |
| 84 | + data = target.read_bytes() |
| 85 | + byte_order, records = _records(data) |
59 | 86 | terminator = next(record for record in records if record.record_type == 999) |
60 | | - insertion = extension.start if extension else terminator.start |
61 | | - target.write_bytes(data[:insertion] + replacement + data[insertion:]) |
| 87 | + long_names = next( |
| 88 | + ( |
| 89 | + record for record in records |
| 90 | + if record.record_type == 7 and _int(data, record.start + 4, byte_order) == 13 |
| 91 | + ), |
| 92 | + None, |
| 93 | + ) |
| 94 | + if long_names is None: |
| 95 | + raise RawDictionaryError("SAV dictionary has no long-variable-name record.") |
| 96 | + payload_start = long_names.start + 16 |
| 97 | + payload = data[payload_start : long_names.end] |
| 98 | + pairs = _long_name_pairs(payload, encoding) |
| 99 | + replacements: dict[str, str] = {} |
| 100 | + updated_pairs: list[tuple[str, str]] = [] |
| 101 | + unresolved = set(names) |
| 102 | + for short_name, long_name in pairs: |
| 103 | + replacement = names.get(long_name) |
| 104 | + if replacement is None: |
| 105 | + updated_pairs.append((short_name, long_name)) |
| 106 | + continue |
| 107 | + replacements[short_name] = _validated_compatible_name(replacement) |
| 108 | + updated_pairs.append((replacements[short_name], long_name)) |
| 109 | + unresolved.remove(long_name) |
| 110 | + if unresolved: |
| 111 | + missing = ", ".join(sorted(unresolved)) |
| 112 | + raise RawDictionaryError( |
| 113 | + "Compatible-name update requires a long-name record for: " + missing |
| 114 | + ) |
| 115 | + |
| 116 | + mutable = bytearray(data) |
| 117 | + for record in records: |
| 118 | + if record.record_type != 2: |
| 119 | + continue |
| 120 | + short_name = bytes(mutable[record.start + 24 : record.start + 32]).decode(encoding).rstrip(" ") |
| 121 | + replacement = replacements.get(short_name) |
| 122 | + if replacement is not None: |
| 123 | + mutable[record.start + 24 : record.start + 32] = replacement.encode("ascii").ljust(8, b" ") |
| 124 | + |
| 125 | + new_payload = b" ".join( |
| 126 | + short_name.encode("ascii") + b"=" + long_name.encode(encoding) |
| 127 | + for short_name, long_name in updated_pairs |
| 128 | + ) |
| 129 | + header = bytearray(mutable[long_names.start : long_names.start + 16]) |
| 130 | + header[12:16] = _pack(len(new_payload), byte_order) |
| 131 | + updated = ( |
| 132 | + bytes(mutable[: long_names.start]) + bytes(header) + new_payload |
| 133 | + + bytes(mutable[long_names.end :]) |
| 134 | + ) |
| 135 | + target.write_bytes(_shift_zsav_offsets( |
| 136 | + updated, |
| 137 | + original_data_start=terminator.end, |
| 138 | + delta=len(updated) - len(data), |
| 139 | + byte_order=byte_order, |
| 140 | + )) |
| 141 | + |
| 142 | + |
| 143 | +def _long_name_pairs(payload: bytes, encoding: str) -> list[tuple[str, str]]: |
| 144 | + pairs: list[tuple[str, str]] = [] |
| 145 | + for raw_pair in payload.split(b" "): |
| 146 | + try: |
| 147 | + raw_short, raw_long = raw_pair.split(b"=", maxsplit=1) |
| 148 | + short_name = raw_short.decode("ascii") |
| 149 | + long_name = raw_long.decode(encoding) |
| 150 | + except (UnicodeDecodeError, ValueError) as error: |
| 151 | + raise RawDictionaryError("Invalid SAV long-variable-name record.") from error |
| 152 | + pairs.append((short_name, long_name)) |
| 153 | + return pairs |
| 154 | + |
| 155 | + |
| 156 | +def _validated_compatible_name(value: str) -> str: |
| 157 | + encoded = str(value).encode("ascii") |
| 158 | + if not 1 <= len(encoded) <= 8: |
| 159 | + raise RawDictionaryError("SPSS compatible variable names must contain one to eight ASCII bytes.") |
| 160 | + first = encoded[:1] |
| 161 | + rest = encoded[1:] |
| 162 | + if not (first.isalpha() or first == b"@"): |
| 163 | + raise RawDictionaryError("SPSS compatible variable names must start with an ASCII letter or @.") |
| 164 | + if any(character not in b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_.$#" |
| 165 | + for character in rest): |
| 166 | + raise RawDictionaryError("SPSS compatible variable name contains an unsupported character.") |
| 167 | + return encoded.decode("ascii") |
| 168 | + |
| 169 | + |
| 170 | +def _shift_zsav_offsets( |
| 171 | + data: bytes, *, original_data_start: int, delta: int, byte_order: str, |
| 172 | +) -> bytes: |
| 173 | + """Shift all absolute ZSAV data offsets after a dictionary-size rewrite.""" |
| 174 | + if not delta or data[:4] != b"$FL3": |
| 175 | + return data |
| 176 | + adjusted = bytearray(data) |
| 177 | + data_start = original_data_start + delta |
| 178 | + _need(adjusted, data_start, 24) |
| 179 | + original_zheader = _int64(adjusted, data_start, byte_order) |
| 180 | + original_ztrailer = _int64(adjusted, data_start + 8, byte_order) |
| 181 | + trailer_length = _int64(adjusted, data_start + 16, byte_order) |
| 182 | + if original_zheader != original_data_start: |
| 183 | + raise RawDictionaryError("Unexpected ZSAV data-header offset.") |
| 184 | + _set_int64(adjusted, data_start, original_zheader + delta, byte_order) |
| 185 | + _set_int64(adjusted, data_start + 8, original_ztrailer + delta, byte_order) |
| 186 | + trailer_start = original_ztrailer + delta |
| 187 | + if trailer_length < 24: |
| 188 | + raise RawDictionaryError("Invalid ZSAV data-trailer length.") |
| 189 | + _need(adjusted, trailer_start, 24) |
| 190 | + block_count = _int(adjusted, trailer_start + 20, byte_order) |
| 191 | + if block_count < 0 or trailer_length != 24 + (block_count * 24): |
| 192 | + raise RawDictionaryError("Invalid ZSAV data-trailer block metadata.") |
| 193 | + _need(adjusted, trailer_start + 24, block_count * 24) |
| 194 | + for index in range(block_count): |
| 195 | + descriptor = trailer_start + 24 + (index * 24) |
| 196 | + _set_int64(adjusted, descriptor, _int64(adjusted, descriptor, byte_order) + delta, byte_order) |
| 197 | + _set_int64( |
| 198 | + adjusted, descriptor + 8, |
| 199 | + _int64(adjusted, descriptor + 8, byte_order) + delta, |
| 200 | + byte_order, |
| 201 | + ) |
| 202 | + return bytes(adjusted) |
62 | 203 |
|
63 | 204 |
|
64 | 205 | def _document_record(lines: Iterable[str], *, encoding: str, byte_order: str) -> bytes: |
@@ -160,6 +301,14 @@ def _int(data: bytes, offset: int, byte_order: str) -> int: |
160 | 301 | return int.from_bytes(data[offset : offset + 4], byte_order, signed=True) |
161 | 302 |
|
162 | 303 |
|
| 304 | +def _int64(data: bytes, offset: int, byte_order: str) -> int: |
| 305 | + return int.from_bytes(data[offset : offset + 8], byte_order, signed=True) |
| 306 | + |
| 307 | + |
| 308 | +def _set_int64(data: bytearray, offset: int, value: int, byte_order: str) -> None: |
| 309 | + data[offset : offset + 8] = int(value).to_bytes(8, byte_order, signed=True) |
| 310 | + |
| 311 | + |
163 | 312 | def _round4(value: int) -> int: |
164 | 313 | return (value + 3) & ~3 |
165 | 314 |
|
|
0 commit comments