Mercurial > hg > cc > cirrus_work
annotate bin/warc.py @ 45:212da3fe3b19
make test 1 idempotent
| author | Henry S. Thompson <ht@inf.ed.ac.uk> |
|---|---|
| date | Wed, 05 Jul 2023 19:32:02 +0100 |
| parents | 69be1131bcc5 |
| children | 44d3a4f4ea51 |
| rev | line source |
|---|---|
| 39 | 1 #!/usr/bin/env python3 |
|
42
689a0e311cd2
make warc.py a library, separate out testing
Henry S. Thompson <ht@inf.ed.ac.uk>
parents:
39
diff
changeset
|
2 '''Stream a gzipped warc format file, invoking a callback on each record. |
| 39 | 3 Callback can be limited by WARC-Type, record part''' |
| 4 | |
|
42
689a0e311cd2
make warc.py a library, separate out testing
Henry S. Thompson <ht@inf.ed.ac.uk>
parents:
39
diff
changeset
|
5 import sys,io |
|
689a0e311cd2
make warc.py a library, separate out testing
Henry S. Thompson <ht@inf.ed.ac.uk>
parents:
39
diff
changeset
|
6 from isal import igzip |
| 39 | 7 |
|
42
689a0e311cd2
make warc.py a library, separate out testing
Henry S. Thompson <ht@inf.ed.ac.uk>
parents:
39
diff
changeset
|
8 def warc(filename,callback,types=['response'],whole=False,parts=7,debug=False): |
| 39 | 9 types=[(t if isinstance(t,bytes) else bytes(t,'utf8')) for t in types] |
| 10 nb=0 | |
|
42
689a0e311cd2
make warc.py a library, separate out testing
Henry S. Thompson <ht@inf.ed.ac.uk>
parents:
39
diff
changeset
|
11 stream=open(filename,'rb',0) |
| 39 | 12 bufsize=128*1024*1024 |
| 13 buf=bytearray(128*1024*1024) | |
| 14 l=b'\r\n' | |
| 15 while not stream.closed: | |
| 16 bp=0 | |
| 17 while l==b'\r\n': | |
| 18 l=stream.readline() | |
| 19 nb+=(ln:=len(l)) | |
| 43 | 20 if ln==0: |
| 21 break | |
| 39 | 22 if l!=b'WARC/1.0\r\n': |
| 23 raise ValueError("Not a WARC file? At %s: %s[%s]"%(nb-len(l), | |
| 24 l.decode('latin-1'),len(l))) | |
| 25 wtype=None | |
| 26 length=None | |
| 27 state=1 | |
| 28 tr=None # Was this record truncated? | |
| 29 while l!=b'\r\n': | |
| 30 # WARC header | |
| 31 if parts & 1: | |
| 32 buf[bp:(bp:=bp+ln)]=l | |
| 33 l=stream.readline() | |
| 34 nb+=(ln:=len(l)) | |
| 35 if l.startswith(b"Content-Length: "): | |
| 36 length=wl=int(l[16:].rstrip()) | |
| 37 elif l.startswith(b"WARC-Truncated: "): | |
| 38 tr=l[16:].rstrip() | |
| 39 tr="EMPTY" if tr=="" else tr | |
| 40 elif l.startswith(b'WARC-Type: '): | |
| 41 wtype = l[11:-2] | |
| 42 start_2=bp | |
| 43 if (wtype in types): | |
| 44 if whole: | |
| 45 buf[bp:(bp:=bp+ln)]=l | |
| 46 elif (parts & 1): | |
| 47 callback(wtype,buf[:start_2],1) | |
| 48 if parts==1: | |
| 49 start_2=0 | |
| 50 else: | |
| 51 start_2=bp | |
| 52 else: | |
| 53 start_2=0 | |
| 54 bv=memoryview(buf)[start_2:start_2+length] | |
| 55 ii=0 | |
| 56 while True and not stream.closed: | |
| 57 if (i:=stream.readinto(bv))==0: | |
| 58 break | |
| 59 ii+=i | |
| 60 if ii>=length: | |
| 61 break | |
| 62 bv=memoryview(buf)[start_2+ii:start_2+length] | |
| 63 if ii!=length: | |
| 64 raise ValueError("Chunk read losing, from %s got %s expected %s"%(nb,ii,length)) | |
| 65 nb+=length | |
| 66 if wtype in types: | |
| 67 if whole: | |
| 68 callback(wtype,buf[0:start_2+length],7) | |
| 69 continue | |
| 70 # Only output parts (1 = WARC header, 2 = HTTP header, 4 = body) that are wanted | |
| 71 bl=None # for HTTP Content-Length for the length of the body? | |
| 72 L_start=start_2 | |
| 73 state=2 | |
| 74 bv=memoryview(buf)[start_2:start_2+length] | |
| 75 with io.BytesIO(bv) as rec_text: | |
| 76 for L in rec_text: | |
| 77 if state==2: | |
| 78 # HTTP header | |
| 79 wl -= len(L) | |
| 80 if not (L==b"" or L.startswith(b"\r")): | |
| 81 # Non-empty, it's (a continuation of) a header | |
| 82 if bl is None and L.startswith(b"Content-Length: "): | |
| 83 bl=int(L[16:].rstrip()) | |
| 84 else: | |
| 85 # Blank line, HTTP header is finished | |
| 86 if parts & 2: | |
| 87 callback(wtype,buf[start_2:start_2+L_start],2) | |
| 88 state=4 | |
| 89 # The above is just for sanity, because we do _not_ | |
| 90 # continue with the outer loop, | |
| 91 # since we can now block-output the entire rest of the | |
| 92 # input buffer. | |
| 93 if bl is not None: | |
| 94 if bl!=wl: | |
| 95 print("length mismatch: %s %s %s here: %s given: %s trunc: %s"%\ | |
| 96 (length,offset,filename,wl,bl,tr),file=sys.stderr) | |
| 97 # HTTP body | |
| 98 balance=start_2+rec_text.tell() | |
| 99 #print(balance,bl,wl,ll,ll-balance,file=sys.stderr) | |
| 100 # Output whatever is left | |
| 101 if parts & 4: | |
| 102 callback(wtype,buf[balance:balance+wl],4) | |
| 103 state=1 | |
| 104 | |
| 105 L_start=rec_text.tell() |