warcforhumans 1.0.0a2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,60 @@
1
+ Metadata-Version: 2.3
2
+ Name: warcforhumans
3
+ Version: 1.0.0a2
4
+ Summary:
5
+ Requires-Dist: zstandard
6
+ Requires-Dist: h11>=0.16.0,<0.17.0
7
+ Requires-Dist: requests>=2.32.5,<3.0.0
8
+ Requires-Dist: typing-extensions>=4.15.0,<5.0.0
9
+ Requires-Python: >=3.12
10
+ Description-Content-Type: text/markdown
11
+
12
+ # warcforhumans
13
+ warcforhumans is a Python package that lets you write WARC records with a simple `requests` session.
14
+
15
+ > [!CAUTION]
16
+ > This project is still a work-in-progress, and has not been extensively tested.
17
+
18
+ # Usage
19
+
20
+ Simple example:
21
+
22
+ ```python
23
+ from warcforhumans.api import WARCWriter
24
+
25
+ warc_writer = WARCWriter("example", rotate_mb=0)
26
+
27
+ s = warc_writer.get_session()
28
+ r = s.get("http://digitaldragon.dev")
29
+ warc_writer.close()
30
+
31
+ ```
32
+
33
+ With ZSTD compression:
34
+
35
+ ```python
36
+ from warcforhumans.api import WARCWriter
37
+ from warcforhumans.compression import ZSTDCompressor
38
+
39
+ warc_writer = WARCWriter("example-$date-$number-$serial",
40
+ compressor=ZSTDCompressor(level=11),
41
+ warcinfo_fields={"operator": "some person"},
42
+ software="example-script/0.1")
43
+ ```
44
+
45
+ or a ZSTD dictionary:
46
+
47
+ ```python
48
+ from warcforhumans.api import WARCWriter
49
+ from warcforhumans.compression import ZSTDCompressor
50
+
51
+ with open("dictionary.zstdict", "rb") as f:
52
+ dictionary = f.read()
53
+
54
+ warc_writer = WARCWriter("example-$date-$number-$serial",
55
+ compressor=ZSTDCompressor(level=11, dictionary=dictionary),
56
+ warcinfo_fields={"operator": "some person"},
57
+ software="example-script/0.1")
58
+ ```
59
+
60
+
@@ -0,0 +1,49 @@
1
+ # warcforhumans
2
+ warcforhumans is a Python package that lets you write WARC records with a simple `requests` session.
3
+
4
+ > [!CAUTION]
5
+ > This project is still a work-in-progress, and has not been extensively tested.
6
+
7
+ # Usage
8
+
9
+ Simple example:
10
+
11
+ ```python
12
+ from warcforhumans.api import WARCWriter
13
+
14
+ warc_writer = WARCWriter("example", rotate_mb=0)
15
+
16
+ s = warc_writer.get_session()
17
+ r = s.get("http://digitaldragon.dev")
18
+ warc_writer.close()
19
+
20
+ ```
21
+
22
+ With ZSTD compression:
23
+
24
+ ```python
25
+ from warcforhumans.api import WARCWriter
26
+ from warcforhumans.compression import ZSTDCompressor
27
+
28
+ warc_writer = WARCWriter("example-$date-$number-$serial",
29
+ compressor=ZSTDCompressor(level=11),
30
+ warcinfo_fields={"operator": "some person"},
31
+ software="example-script/0.1")
32
+ ```
33
+
34
+ or a ZSTD dictionary:
35
+
36
+ ```python
37
+ from warcforhumans.api import WARCWriter
38
+ from warcforhumans.compression import ZSTDCompressor
39
+
40
+ with open("dictionary.zstdict", "rb") as f:
41
+ dictionary = f.read()
42
+
43
+ warc_writer = WARCWriter("example-$date-$number-$serial",
44
+ compressor=ZSTDCompressor(level=11, dictionary=dictionary),
45
+ warcinfo_fields={"operator": "some person"},
46
+ software="example-script/0.1")
47
+ ```
48
+
49
+
@@ -0,0 +1,32 @@
1
+ [build-system]
2
+ requires = ["uv_build>=0.12.17,<1.0.0"]
3
+ build-backend = "uv_build"
4
+
5
+ [project]
6
+ name = "warcforhumans"
7
+ version = "1.0.0a2"
8
+ description = ""
9
+ authors = []
10
+ readme = "README.md"
11
+ requires-python = ">=3.12"
12
+ dependencies = [
13
+ "zstandard",
14
+ "h11 (>=0.16.0,<0.17.0)",
15
+ "requests (>=2.32.5,<3.0.0)",
16
+ "typing-extensions (>=4.15.0,<5.0.0)",
17
+ ]
18
+
19
+ [dependency-groups]
20
+ dev = ["pytest"]
21
+
22
+ [tool.pytest]
23
+ minversion = "9.0"
24
+ addopts = [
25
+ "-ra",
26
+ "-q",
27
+ ]
28
+ testpaths = ["tests"]
29
+
30
+ [tool.uv.build-backend]
31
+ module-name = ["warcforhumans"]
32
+ module-root = ""
@@ -0,0 +1,34 @@
1
+ [build-system]
2
+ requires = ["uv_build>=0.12.17,<1.0.0"]
3
+ build-backend = "uv_build"
4
+
5
+
6
+ [project]
7
+ name = "warcforhumans"
8
+ version = "1.0.0a2"
9
+ description = ""
10
+ authors = []
11
+ readme = "README.md"
12
+ requires-python = ">=3.12"
13
+ dependencies = [
14
+ "zstandard",
15
+ "h11 (>=0.16.0,<0.17.0)",
16
+ "requests (>=2.32.5,<3.0.0)",
17
+ "typing-extensions (>=4.15.0,<5.0.0)",
18
+ ]
19
+
20
+ [dependency-groups]
21
+ dev = [
22
+ "pytest"
23
+ ]
24
+
25
+ [tool.pytest]
26
+ minversion = "9.0"
27
+ addopts = ["-ra", "-q"]
28
+ testpaths = [
29
+ "tests"
30
+ ]
31
+
32
+ [tool.uv.build-backend]
33
+ module-name = ["warcforhumans"]
34
+ module-root = ""
File without changes
@@ -0,0 +1,506 @@
1
+ from __future__ import annotations
2
+
3
+ import base64
4
+ import hashlib
5
+ import io
6
+ import random
7
+ import string
8
+ import sys
9
+ import tempfile
10
+ import uuid
11
+ from datetime import datetime, timezone
12
+ from importlib.metadata import version
13
+ from io import BufferedRandom
14
+ from socket import socket
15
+ from ssl import SSLSocket
16
+ from typing import Iterator, BinaryIO
17
+
18
+ import requests
19
+ from _hashlib import HASH
20
+
21
+ from warcforhumans.compression import Compressor
22
+
23
+
24
+ class WARCRecord:
25
+ WARC_RECORD_ID = "WARC-Record-ID"
26
+ CONTENT_LENGTH = "Content-Length"
27
+ WARC_DATE = "WARC-Date"
28
+ WARC_TYPE = "WARC-Type"
29
+ CONTENT_TYPE = "Content-Type"
30
+ CONTENT_HTTP_REQUEST = "application/http;msgtype=request"
31
+ CONTENT_HTTP_RESPONSE = "application/http;msgtype=response"
32
+ WARC_CONCURRENT_TO = "WARC-Concurrent-To"
33
+ WARC_BLOCK_DIGEST = "WARC-Block-Digest"
34
+ WARC_PAYLOAD_DIGEST = "WARC-Payload-Digest"
35
+ WARC_IP_ADDRESS = "WARC-IP-Address"
36
+ WARC_REFERS_TO = "WARC-Refers-To"
37
+ WARC_REFERS_TO_TARGET_URI = "WARC-Refers-To-Target-URI"
38
+ WARC_REFERS_TO_DATE = "WARC-Refers-To-Date"
39
+ WARC_TARGET_URI = "WARC-Target-URI"
40
+ WARC_TRUNCATED = "WARC-Truncated"
41
+ WARC_WARCINFO_ID = "WARC-Warcinfo-ID"
42
+ WARC_FILENAME = "WARC-Filename"
43
+ WARC_PROFILE = "WARC-Profile"
44
+ WARC_IDENTIFIED_PAYLOAD_TYPE = "WARC-Identified-Payload-Type"
45
+ WARC_SEGMENT_NUMBER = "WARC-Segment-Number"
46
+ WARC_SEGMENT_ORIGIN_ID = "WARC-Segment-Origin-ID"
47
+ WARC_SEGMENT_TOTAL_LENGTH = "WARC-Segment-Total-Length"
48
+ WARC_PROTOCOL = "WARC-Protocol" # https://github.com/iipc/warc-specifications/issues/42
49
+ WARC_CIPHER_SUITE = "WARC-Cipher-Suite" # https://github.com/iipc/warc-specifications/issues/86
50
+ HTTP_1_1 = "http/1.1"
51
+
52
+ def __init__(self, record_type: str = None, content_type: str = None, url: str | None = None, sock: socket = None):
53
+ """
54
+ Creates a new WARC record with a random ID.
55
+ :param record_type: The type of WARC record. Equivalent to ``set_type``
56
+ :param content_type: The ``Content-Type`` of the WARC record
57
+ :param url: The URL of the WARC record
58
+ :param sock: Add headers based on the socket that this record was captured from. Equivalent to
59
+ ``add_headers_for_socket``
60
+ """
61
+ # https://iipc.github.io/warc-specifications/specifications/warc-format/warc-1.1-annotated/#warc-type-mandatory
62
+
63
+ self.headers: dict[str, list[str]] = {}
64
+ self.content = None
65
+ self._close_content_stream : bool = False
66
+
67
+ self.get_id()
68
+
69
+ if record_type is not None:
70
+ self.set_type(record_type)
71
+
72
+ if content_type is not None:
73
+ self.set_header(WARCRecord.CONTENT_TYPE, content_type)
74
+
75
+ if url is not None:
76
+ self.set_header(WARCRecord.WARC_TARGET_URI, url)
77
+
78
+ if sock is not None:
79
+ self.add_headers_for_socket(sock)
80
+
81
+ self._partial_content: bytes | BufferedRandom = b""
82
+ self.max_in_memory_record_size: int = 1024 * 1024
83
+
84
+ def set_header(self, key: str, value: str | list[str]) -> None:
85
+ """
86
+ Sets a header on the record, **overwriting** the value if one is already set.
87
+ :param key: Header name
88
+ :param value: Header value
89
+ :return:
90
+ """
91
+ if isinstance(value, list):
92
+ self.headers[key] = value
93
+ else:
94
+ self.headers[key] = [value]
95
+
96
+ def add_header(self, key: str, value: str):
97
+ """
98
+ Adds a header to the record, **repeating the header** if the header is already set.
99
+ :param key: Header name
100
+ :param value: Header value
101
+ :return:
102
+ """
103
+ if key in self.headers:
104
+ self.headers[key].append(value)
105
+ else:
106
+ self.headers[key] = [value]
107
+
108
+ def add_headers(self, headers: dict[str, list[str]]):
109
+ self.headers.update(headers)
110
+
111
+ def set_type(self, record_type: str):
112
+ valid_warc_record_types = {"warcinfo", "response", "resource", "request", "metadata", "revisit", "conversion",
113
+ "continuation"}
114
+ if record_type not in valid_warc_record_types:
115
+ raise ValueError(f"Invalid WARC record type: {record_type}")
116
+
117
+ self.set_header(WARCRecord.WARC_TYPE, record_type)
118
+
119
+ def set_content(self, content: bytes | BinaryIO, content_type: str | None = None, block_digest: str | None = None, close: bool = False) -> None:
120
+ """
121
+ Sets the body content of the WARC record.
122
+ :param content: The content of the record, as bytes or a seekable file object.
123
+ :param content_type: The ``Content-Type`` to set for the record.
124
+ :param block_digest: A hash of content (as a hash object), used as the ``WARC-Block-Digest``. If not set, one
125
+ will be generated with SHA-1.
126
+ :param close: If content is a file object, whether it should be closed after the record is written.
127
+ :return:
128
+ """
129
+ if content_type:
130
+ self.set_header(WARCRecord.CONTENT_TYPE, content_type)
131
+
132
+ if isinstance(content, bytes):
133
+ self._set_content_bytes(content, block_digest=block_digest)
134
+ elif hasattr(content, "read") and hasattr(content, "seek"):
135
+ self._set_content_stream(content, close=close, block_digest=block_digest)
136
+ else:
137
+ raise ValueError("Unknown record content - expected a bytes or BufferedRandom like object")
138
+
139
+
140
+ def _set_content_bytes(self, content: bytes, block_digest = None) -> None:
141
+ self.content = content
142
+ self.set_header(WARCRecord.CONTENT_LENGTH, str(len(content)))
143
+
144
+ if not block_digest:
145
+ block_digest = hashlib.sha1(content)
146
+ self.set_header(WARCRecord.WARC_BLOCK_DIGEST, hash_to_string(block_digest))
147
+
148
+
149
+ def _set_content_stream(self, stream: BinaryIO, close: bool = False, block_digest = None) -> None:
150
+ self.content = stream
151
+ stream.seek(0, io.SEEK_END)
152
+ self.set_header(WARCRecord.CONTENT_LENGTH, str((stream.tell())))
153
+ self._close_content_stream = close
154
+
155
+ if not block_digest:
156
+ block_digest = hashlib.sha1()
157
+ stream.seek(0)
158
+ while chunk := stream.read(2048):
159
+ block_digest.update(chunk)
160
+
161
+ self.set_header(WARCRecord.WARC_BLOCK_DIGEST, hash_to_string(block_digest))
162
+
163
+ def partial_content(self, content: bytes, finish: bool = False) -> None:
164
+ if isinstance(self._partial_content, bytes):
165
+ if len(self._partial_content) + len(content) > self.max_in_memory_record_size:
166
+ file = tempfile.TemporaryFile()
167
+ _ = file.write(self._partial_content)
168
+ _ = file.write(content)
169
+ self._partial_content = file
170
+ else:
171
+ self._partial_content += content
172
+ else:
173
+ _ = self._partial_content.write(content)
174
+
175
+ if finish:
176
+ self.set_content(self._partial_content, close=True)
177
+
178
+
179
+
180
+ def date(self, date: datetime | None = None):
181
+ """
182
+ Sets the timestamp of the WARC to the datetime provided, or the current time if no time is provided.
183
+ :param date: The datetime to set.
184
+ :return:
185
+ """
186
+ templ = "%Y-%m-%dT%H:%M:%SZ"
187
+ if date is not None:
188
+ if date.tzinfo is None or date.tzinfo != timezone.utc:
189
+ raise ValueError("Timezone for the WARC-Date must be in UTC.")
190
+
191
+ self.set_header(WARCRecord.WARC_DATE, date.strftime(templ))
192
+ else:
193
+ self.set_header(WARCRecord.WARC_DATE, datetime.now(timezone.utc).strftime(templ))
194
+
195
+ def get_id(self) -> str:
196
+ """
197
+ Gets the ID of the record. If there isn't one already, one will be randomly generated.
198
+ :return:
199
+ """
200
+ if WARCRecord.WARC_RECORD_ID not in self.headers:
201
+ self.set_header(WARCRecord.WARC_RECORD_ID, f"<{str(uuid.uuid4().urn)}>")
202
+ return self.headers[WARCRecord.WARC_RECORD_ID][0]
203
+
204
+ def get_type(self) -> str | None:
205
+ if WARCRecord.WARC_TYPE not in self.headers:
206
+ return None
207
+ return self.headers[WARCRecord.WARC_TYPE][0]
208
+
209
+ def add_headers_for_socket(self, sock: socket):
210
+ """
211
+ Sets headers of the WARC record based on the given socket. ``WARC-IP-Address`` will be set for all sockets.
212
+ For an SSLSocket,``WARC-Protocol`` and ``WARC-Cipher-Suite`` will also be set to document the encryption used.
213
+ :param sock: The socket that this record was captured from
214
+ :return:
215
+ """
216
+ self.set_header(WARCRecord.WARC_IP_ADDRESS, sock.getpeername()[0])
217
+
218
+ if isinstance(sock, SSLSocket):
219
+ encryption_protocol, protocol_version = sock.cipher()[1].split("v")
220
+ self.add_header(WARCRecord.WARC_PROTOCOL, encryption_protocol.lower() + "/" + protocol_version)
221
+ self.add_header(WARCRecord.WARC_CIPHER_SUITE, sock.cipher()[0])
222
+
223
+ def concurrent(self, record: WARCRecord) -> None:
224
+ """
225
+ Sets ``WARC-Concurrent-To`` headers on both given records, pointing at each other.
226
+ :param record: The other record.
227
+ :return:
228
+ """
229
+ if self.headers[WARCRecord.WARC_DATE] != record.headers[WARCRecord.WARC_DATE]:
230
+ raise ValueError("The WARC-Date fields of records that should be concurrent are not the same.")
231
+
232
+ self.set_header(WARCRecord.WARC_CONCURRENT_TO, record.get_id())
233
+ record.set_header(WARCRecord.WARC_CONCURRENT_TO, self.get_id())
234
+
235
+ def serialize_stream(self) -> Iterator[bytes]:
236
+ """
237
+ Serializes this WARC record as bytes and yields it in chunks.
238
+ :return: Iterator[bytes] of the serialized record chunks
239
+ """
240
+
241
+ if self.content is None:
242
+ raise ValueError("Content is not set")
243
+
244
+ mandatory_headers = [WARCRecord.WARC_RECORD_ID, WARCRecord.CONTENT_LENGTH, WARCRecord.WARC_DATE, WARCRecord.WARC_TYPE]
245
+ for header in mandatory_headers:
246
+ if header not in self.headers:
247
+ raise ValueError(f"Mandatory header {header} is missing")
248
+
249
+ yield b"WARC/1.1\r\n"
250
+ for key, value in self.headers.items():
251
+ for v in value:
252
+ yield f"{key}: {v}\r\n".encode("utf-8")
253
+ yield b"\r\n"
254
+
255
+ if isinstance(self.content, bytes):
256
+ yield self.content
257
+ else:
258
+ chunk_size = 8192 # todo we should use a set chunk size everywhere
259
+ self.content.seek(0)
260
+ while True:
261
+ chunk = self.content.read(chunk_size)
262
+ if not chunk:
263
+ break
264
+ yield chunk
265
+
266
+ yield b"\r\n\r\n"
267
+
268
+ def close(self):
269
+ if self._close_content_stream and hasattr(self.content, 'close'):
270
+ self.content.close()
271
+
272
+ class WARCFile:
273
+ def __init__(self, file_path: str, create_warcinfo: bool = True, warcinfo_fields = None, compressor: Compressor = None, software: str = ""):
274
+ """
275
+ Creates a new WARCFile to write WARC records to.
276
+ :param file_path: Where the file should go, without ``.warc``
277
+ :param create_warcinfo: Whether to create a warcinfo record for this file
278
+ :param warcinfo_fields: Fields to add to the warcinfo record.
279
+ :param compressor: ``Compressor`` to use for the file. The WARC will be uncompressed if not specified.
280
+ :param software: The name of the software to add to warcinfo. warcforhumans will add its own name and version
281
+ to the front of the provided string
282
+ """
283
+ self._warcinfo_record = None
284
+ self._pending_records = []
285
+
286
+ if not compressor:
287
+ self._compressor = Compressor()
288
+ else:
289
+ self._compressor = compressor
290
+
291
+ self.file_path = file_path + ".warc" + self._compressor.file_extension()
292
+ self.file = open(self.file_path, "ab")
293
+ self._compressor.start(self.file)
294
+
295
+ if create_warcinfo:
296
+ self.create_warcinfo_record(fields=warcinfo_fields,software=software)
297
+
298
+ def create_warcinfo_record(self, software: str = "", fields: dict[str, str] = None):
299
+ """
300
+ Creates a warcinfo record for the WARC file.
301
+ :param software: The name of the software to add to warcinfo. warcforhumans will add its own name and version
302
+ to the front of the provided string
303
+ :param fields: Fields to add to the warcinfo record.
304
+ :return:
305
+ """
306
+ if fields is None:
307
+ fields = {}
308
+
309
+ warc_record = WARCRecord("warcinfo", "application/warc-fields")
310
+ warc_record.set_header(WARCRecord.WARC_FILENAME, self.file.name)
311
+
312
+ fields["software"] = f"warcforhumans/{version("warcforhumans")} {software}"
313
+ fields["format"] = "WARC File Format 1.1"
314
+ fields["conformsTo"] = "https://bibnum.bnf.fr/WARC/WARC_ISO_28500_version1-1_latestdraft.pdf"
315
+ fields["python-version"] = "python/" + sys.version.replace("\n", "")
316
+
317
+ body = str()
318
+ for key, value in fields.items():
319
+ body += f"{key}: {value}\r\n"
320
+
321
+ warc_record.set_content(body.encode("utf-8"))
322
+ warc_record.date()
323
+ self.write_record(warc_record, write_warcinfo_header=False)
324
+ self._warcinfo_record = warc_record
325
+
326
+ def write_record(self, record: WARCRecord, write_warcinfo_header: bool = True):
327
+ """
328
+ Write a record to the file.
329
+ :param record: The record to write
330
+ :param write_warcinfo_header: Whether to write a ``WARC-Warcinfo-ID`` header for this record (if a warcinfo
331
+ record for this file exists).
332
+ :return:
333
+ """
334
+ if self._warcinfo_record is not None and write_warcinfo_header:
335
+ record.set_header(WARCRecord.WARC_WARCINFO_ID, self._warcinfo_record.get_id())
336
+
337
+ self._compressor.write_record(record, self.file)
338
+
339
+ record.close()
340
+ self.file.flush()
341
+
342
+ def close(self):
343
+ self.file.close()
344
+
345
+ class WARCWriter:
346
+ def __init__(self,
347
+ template: str,
348
+ compressor: Compressor = None,
349
+ rotate_mb: int = 15*1024,
350
+ software: str = "",
351
+ warcinfo_fields: dict[str, str] = None,
352
+ revisit: bool = True
353
+ ):
354
+ """
355
+ Creates a WARCWriter, to manage writing WARC records.
356
+ :param template: Name template for the created WARC files. Available variables:
357
+ ``$date`` - the datetime the file was started at (``yyyymmddHHMMSS``)
358
+ ``$number`` - the number of files made by this writer, starting from 00000
359
+ ``$serial`` - 8 character random string
360
+ :param compressor: ``Compressor`` to use for the file. The WARC will be uncompressed if not specified.
361
+ :param rotate_mb: Automatically start a new file when the old one gets above this size. Set to -1 to never start
362
+ a new file.
363
+ :param software: The name of the software to add to warcinfo. warcforhumans will add its own name and version
364
+ to the front of the provided string
365
+ :param warcinfo_fields: Fields to add to the warcinfo record.
366
+ :param revisit: Whether to keep track of the information necessary (in memory) for ``check_for_revisit``
367
+ """
368
+ self.warc_file = None
369
+ self.compressor = compressor if compressor else Compressor()
370
+ self.template = template
371
+ self.rotate_mb = rotate_mb
372
+ self.pending_records: list[WARCRecord] = []
373
+ self.software = software
374
+ self.warcinfo_fields = warcinfo_fields
375
+ self.files_made = 0
376
+ self.closed = False
377
+ self.revisit = revisit
378
+ if revisit:
379
+ self.revisit_cache: dict[str, tuple[str, str, str]] = {} # payload-digest -> (warc-record-id, warc-date, warc-target-uri)
380
+
381
+ def _create_file(self, rotate = True):
382
+ if ( self.warc_file
383
+ and rotate
384
+ and self.rotate_mb > 0
385
+ and self.warc_file.file.tell() >= self.rotate_mb * 1024 * 1024
386
+ ):
387
+ self.warc_file.close()
388
+ self.warc_file = None
389
+ if not self.warc_file:
390
+ template_data = {"date": datetime.now(timezone.utc).strftime("%Y%m%d%H%M%S"),
391
+ "number": f"{self.files_made:05d}",
392
+ "serial": ''.join(random.choices(string.ascii_lowercase + string.digits, k=8))}
393
+ name = string.Template(self.template).substitute(template_data)
394
+ self.warc_file = WARCFile(name, compressor=self.compressor, software=self.software, warcinfo_fields=self.warcinfo_fields)
395
+ self.files_made += 1
396
+
397
+ def flush_pending(self):
398
+ """
399
+ Write records that have been finished but not yet written.
400
+ :return:
401
+ """
402
+ temp = self.pending_records
403
+ self.pending_records = []
404
+ self.write_records(temp)
405
+
406
+ def discard_pending(self):
407
+ """
408
+ Discard records that have been finished but not yet written.
409
+ :return:
410
+ """
411
+ self.pending_records = []
412
+
413
+ def discard(self, record_id):
414
+ """
415
+ Discard a specific pending record by its ID
416
+ :param record_id: ID of the record to discard
417
+ :return:
418
+ """
419
+ self.pending_records = [r for r in self.pending_records if r.get_id() != record_id]
420
+
421
+
422
+ def write_record(self, record: WARCRecord, rotate = True):
423
+ """
424
+ Write a record to the current file.
425
+ :param record: The record to write.
426
+ :param rotate: If disabled, the writer will not start a new file before writing this record, even if it is above
427
+ the rotation size. This can be useful to ensure that request/response pairs don't get split between files.
428
+ :return:
429
+ """
430
+ if self.closed:
431
+ raise ValueError("WARCWriter is closed")
432
+
433
+ self._create_file(rotate=rotate)
434
+ if self.revisit and record.get_type() == "response":
435
+ payload_digest = record.headers.get(WARCRecord.WARC_PAYLOAD_DIGEST, [None])[0]
436
+ date = record.headers.get(WARCRecord.WARC_DATE, [None])[0]
437
+ target_uri = record.headers.get(WARCRecord.WARC_TARGET_URI, [None])[0]
438
+ record_id = record.get_id()
439
+ if payload_digest and date and target_uri:
440
+ self.revisit_cache[payload_digest] = (record_id, date, target_uri)
441
+
442
+ self.warc_file.write_record(record)
443
+
444
+ def check_for_revisit(self, payload_digest: str) -> tuple[bool, dict[str, list[str]]]:
445
+ """
446
+ Checks whether a given payload digest has been seen before by this writer, and returns the necessary WARC
447
+ headers for a revisit if the payload was seen. Will always return false if revisits have been disabled on this
448
+ WARCWriter.
449
+ :param payload_digest: The payload digest to search for
450
+ :return: a tuple (boolean match found, dict of warc headers to add on a revisit)
451
+ """
452
+ if not self.revisit:
453
+ return False, {}
454
+
455
+ if not payload_digest in self.revisit_cache:
456
+ return False, {}
457
+
458
+ record_id, date, target_uri = self.revisit_cache[payload_digest]
459
+ headers = {
460
+ WARCRecord.WARC_REFERS_TO: [record_id],
461
+ WARCRecord.WARC_REFERS_TO_DATE: [date],
462
+ WARCRecord.WARC_REFERS_TO_TARGET_URI: [target_uri]
463
+ }
464
+ return True, headers
465
+
466
+
467
+ def write_records(self, records: list[WARCRecord], rotate_between=False):
468
+ """
469
+ Write multiple records to the WARC file
470
+ :param records: The records to write
471
+ :param rotate_between: Whether the records to be written can be split between different WARC files
472
+ :return:
473
+ """
474
+ for record in records:
475
+ self.write_record(record, rotate=rotate_between)
476
+
477
+ def close(self):
478
+ if not self.closed:
479
+ self.flush_pending()
480
+ if self.warc_file:
481
+ self.warc_file.close()
482
+ self.closed = True
483
+
484
+ def get_session(self) -> requests.Session:
485
+ """
486
+ :return: A requests Session that will write WARC files using this WARCWriter.
487
+ """
488
+ from warcforhumans.capture.adapter import WARCHTTPAdapter
489
+
490
+ s = requests.Session()
491
+ s.mount("http://", WARCHTTPAdapter(warc_writer=self))
492
+ s.mount("https://", WARCHTTPAdapter(warc_writer=self))
493
+ return s
494
+
495
+ def hash_to_string(h: HASH) -> str:
496
+ """
497
+ Converts a hashlib hash object to a string for writing WARC digest headers.
498
+ :param h: The hash to convert
499
+ :return: A string like "sha1:AIKLJM2V2EOKR4WOIWUWRQTEMUN57P4D"
500
+ """
501
+ # Typical encoding is base 16 per https://github.com/iipc/warc-specifications/issues/80#issuecomment-1637084423
502
+ base_16 = {"md5", "sha256", "sha512"}
503
+ if h.name in base_16:
504
+ return f"{h.name}:{h.hexdigest()}"
505
+
506
+ return f"{h.name}:{base64.b32encode(h.digest()).decode("utf-8")}"