|
1 | 1 | """Common mixins""" |
2 | 2 |
|
3 | | -import datetime as dt |
4 | 3 | import inspect |
5 | | -import io |
6 | 4 | import logging |
7 | | -import socket |
8 | | -import time |
9 | 5 | import warnings |
10 | 6 | from abc import ABCMeta, abstractmethod |
11 | | -from collections.abc import Iterator |
12 | | -from typing import BinaryIO, Optional, Protocol, Union |
| 7 | +from typing import Optional, Protocol, Union |
13 | 8 |
|
14 | 9 | import lxml.etree as ElementTree |
15 | 10 | from lxml.builder import ElementMaker |
@@ -387,162 +382,6 @@ def _read_from_binary_as_int(self, nbits: int) -> int: |
387 | 382 | return int_data |
388 | 383 |
|
389 | 384 |
|
390 | | -def fixed_length_generator( |
391 | | - binary_data: Union[BinaryIO, socket.socket, bytes], |
392 | | - *, |
393 | | - packet_length_bytes: int, |
394 | | - buffer_read_size_bytes: Optional[int] = None, |
395 | | - show_progress: bool = False, |
396 | | -) -> Iterator[bytes]: |
397 | | - """A generator that yields fixed-length chunks from binary_data. |
398 | | -
|
399 | | - Parameters |
400 | | - ---------- |
401 | | - binary_data : Union[BinaryIO, socket.socket, bytes] |
402 | | - Binary data source. |
403 | | - packet_length_bytes : int |
404 | | - Number of bytes per packet to yield. |
405 | | - buffer_read_size_bytes : int, optional |
406 | | - Number of bytes to read from the source per read. |
407 | | - show_progress : bool |
408 | | - If True, prints a status bar. |
409 | | -
|
410 | | - Yields |
411 | | - ------ |
412 | | - bytes |
413 | | - Fixed-length packet bytes. |
414 | | - """ |
415 | | - n_bytes_parsed = 0 # Keep track of how many bytes we have parsed |
416 | | - n_packets_parsed = 0 # Keep track of how many packets we have parsed |
417 | | - read_buffer, total_length_bytes, read_bytes_from_source, buffer_read_size_bytes = _setup_binary_reader( |
418 | | - binary_data, buffer_read_size_bytes |
419 | | - ) |
420 | | - current_pos = 0 # Keep track of where we are in the buffer |
421 | | - start_time = time.time_ns() |
422 | | - while True: |
423 | | - if total_length_bytes and n_bytes_parsed == total_length_bytes: |
424 | | - break |
425 | | - if show_progress: |
426 | | - _print_progress( |
427 | | - current_bytes=n_bytes_parsed, |
428 | | - total_bytes=total_length_bytes, |
429 | | - start_time_ns=start_time, |
430 | | - current_packets=n_packets_parsed, |
431 | | - ) |
432 | | - if current_pos > 20_000_000: |
433 | | - # Only trim the buffer after 20 MB read to prevent modifying |
434 | | - # the bitstream and trimming after every packet |
435 | | - read_buffer = read_buffer[current_pos:] |
436 | | - current_pos = 0 |
437 | | - while len(read_buffer) - current_pos < packet_length_bytes: |
438 | | - result = read_bytes_from_source(buffer_read_size_bytes) |
439 | | - if not result: |
440 | | - break |
441 | | - read_buffer += result |
442 | | - packet_bytes = read_buffer[current_pos : current_pos + packet_length_bytes] |
443 | | - current_pos += packet_length_bytes |
444 | | - n_packets_parsed += 1 |
445 | | - n_bytes_parsed += packet_length_bytes |
446 | | - yield packet_bytes |
447 | | - if show_progress: |
448 | | - _print_progress( |
449 | | - current_bytes=n_bytes_parsed, |
450 | | - total_bytes=total_length_bytes, |
451 | | - start_time_ns=start_time, |
452 | | - current_packets=n_packets_parsed, |
453 | | - end="\n", |
454 | | - log=True, |
455 | | - ) |
456 | | - |
457 | | - |
458 | | -def _setup_binary_reader(binary_data, buffer_read_size_bytes=None): |
459 | | - """Helper to set up reading from binary_data (file, socket, bytes). Returns: |
460 | | - read_buffer, total_length_bytes, read_bytes_from_source, buffer_read_size_bytes |
461 | | - """ |
462 | | - read_buffer = b"" |
463 | | - # ======== |
464 | | - # Set up the reader based on the type of binary_data |
465 | | - # ======== |
466 | | - if isinstance(binary_data, io.BufferedIOBase): |
467 | | - if buffer_read_size_bytes is None: |
468 | | - # Default to a full read of the file |
469 | | - buffer_read_size_bytes = -1 |
470 | | - total_length_bytes = binary_data.seek(0, io.SEEK_END) # This is probably preferable to len |
471 | | - binary_data.seek(0, 0) |
472 | | - logger.info( |
473 | | - f"Creating packet generator from a filelike object, {binary_data}. " |
474 | | - f"Total length is {total_length_bytes} bytes" |
475 | | - ) |
476 | | - read_bytes_from_source = binary_data.read |
477 | | - elif isinstance(binary_data, socket.socket): # It's a socket and we don't know how much data we will get |
478 | | - logger.info("Creating packet generator to read from a socket. Total length to parse is unknown.") |
479 | | - total_length_bytes = None # We don't know how long it is |
480 | | - if buffer_read_size_bytes is None: |
481 | | - # Default to 4096 bytes from a socket |
482 | | - buffer_read_size_bytes = 4096 |
483 | | - read_bytes_from_source = binary_data.recv |
484 | | - elif isinstance(binary_data, bytes): |
485 | | - read_buffer = binary_data |
486 | | - total_length_bytes = len(read_buffer) |
487 | | - read_bytes_from_source = None # No data to read, we've filled the read_buffer already |
488 | | - logger.info(f"Creating packet generator from a bytes object. Total length is {total_length_bytes} bytes") |
489 | | - elif isinstance(binary_data, io.TextIOWrapper): |
490 | | - raise OSError("Packet data file opened in TextIO mode. You must open packet data in binary mode.") |
491 | | - else: |
492 | | - raise OSError(f"Unrecognized data source: {binary_data}") |
493 | | - return read_buffer, total_length_bytes, read_bytes_from_source, buffer_read_size_bytes |
494 | | - |
495 | | - |
496 | | -def _print_progress( |
497 | | - *, |
498 | | - current_bytes: int, |
499 | | - total_bytes: Optional[int], |
500 | | - start_time_ns: int, |
501 | | - current_packets: int, |
502 | | - end: str = "\r", |
503 | | - log: bool = False, |
504 | | -): |
505 | | - """Prints a progress bar, including statistics on parsing rate. |
506 | | -
|
507 | | - Parameters |
508 | | - ---------- |
509 | | - current_bytes : int |
510 | | - Number of bytes parsed so far. |
511 | | - total_bytes : Optional[int] |
512 | | - Number of total bytes to parse, if known. None otherwise. |
513 | | - current_packets : int |
514 | | - Number of packets parsed so far. |
515 | | - start_time_ns : int |
516 | | - Start time on system clock, in nanoseconds. |
517 | | - end : str |
518 | | - Print function end string. Default is `\\r` to create a dynamically updating loading bar. |
519 | | - log : bool |
520 | | - If True, log the progress bar at INFO level. |
521 | | - """ |
522 | | - progress_char = "=" |
523 | | - bar_length = 20 |
524 | | - |
525 | | - if total_bytes is not None: # If we actually have an endpoint (i.e. not using a socket) |
526 | | - percentage = int((current_bytes / total_bytes) * 100) # Percent Completed Calculation |
527 | | - progress = int((bar_length * current_bytes) / total_bytes) # Progress Done Calculation |
528 | | - else: |
529 | | - percentage = "???" |
530 | | - progress = 0 |
531 | | - |
532 | | - # Fast calls initially on Windows can result in a zero elapsed time |
533 | | - elapsed_ns = max(time.time_ns() - start_time_ns, 1) |
534 | | - delta = dt.timedelta(microseconds=elapsed_ns / 1e3) |
535 | | - kbps = int(current_bytes * 8e6 / elapsed_ns) # 8 bits per byte, 1E9 s per ns, 1E3 bits per kb |
536 | | - pps = int(current_packets * 1e9 / elapsed_ns) |
537 | | - info_str = ( |
538 | | - f"[Elapsed: {delta}, Parsed {current_bytes} bytes ({current_packets} packets) at {kbps}kb/s ({pps}pkts/s)]" |
539 | | - ) |
540 | | - loadbar = f"Progress: [{progress * progress_char:{bar_length}}]{percentage}% {info_str}" |
541 | | - print(loadbar, end=end) |
542 | | - if log: |
543 | | - logger.info(loadbar) |
544 | | - |
545 | | - |
546 | 385 | def _extract_bits(data: bytes, start_bit: int, nbits: int): |
547 | 386 | """Extract nbits from the data starting from the least significant end. |
548 | 387 |
|
|
0 commit comments