igor/util.py

   1 # Copyright (C) 2012 W. Trevor King <wking@tremily.us>
   2 #
   3 # This file is part of igor.
   4 #
   5 # igor is free software: you can redistribute it and/or modify it under the
   6 # terms of the GNU Lesser General Public License as published by the Free
   7 # Software Foundation, either version 3 of the License, or (at your option) any
   8 # later version.
   9 #
  10 # igor is distributed in the hope that it will be useful, but WITHOUT ANY
  11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  12 # A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
  13 # details.
  14 #
  15 # You should have received a copy of the GNU Lesser General Public License
  16 # along with igor.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 "Utility functions for handling buffers"
  19
  20 import sys as _sys
  21
  22 import numpy as _numpy
  23
  24
  25 def _ord(byte):
  26     r"""Convert a byte to an integer.
  27
  28     >>> buffer = b'\x00\x01\x02'
  29     >>> [_ord(b) for b in buffer]
  30     [0, 1, 2]
  31     """
  32     if _sys.version_info >= (3,):
  33         return byte
  34     else:
  35         return ord(byte)
  36
  37 def hex_bytes(buffer, spaces=None):
  38     r"""Pretty-printing for binary buffers.
  39
  40     >>> hex_bytes(b'\x00\x01\x02\x03\x04')
  41     '0001020304'
  42     >>> hex_bytes(b'\x00\x01\x02\x03\x04', spaces=1)
  43     '00 01 02 03 04'
  44     >>> hex_bytes(b'\x00\x01\x02\x03\x04', spaces=2)
  45     '0001 0203 04'
  46     >>> hex_bytes(b'\x00\x01\x02\x03\x04\x05\x06', spaces=2)
  47     '0001 0203 0405 06'
  48     >>> hex_bytes(b'\x00\x01\x02\x03\x04\x05\x06', spaces=3)
  49     '000102 030405 06'
  50     """
  51     hex_bytes = ['{:02x}'.format(_ord(x)) for x in buffer]
  52     if spaces is None:
  53         return ''.join(hex_bytes)
  54     elif spaces is 1:
  55         return ' '.join(hex_bytes)
  56     for i in range(len(hex_bytes)//spaces):
  57         hex_bytes.insert((spaces+1)*(i+1)-1, ' ')
  58     return ''.join(hex_bytes)
  59
  60 def assert_null(buffer, strict=True):
  61     r"""Ensure an input buffer is entirely zero.
  62
  63     >>> import sys
  64     >>> assert_null(b'')
  65     >>> assert_null(b'\x00\x00')
  66     >>> assert_null(b'\x00\x01\x02\x03')
  67     Traceback (most recent call last):
  68       ...
  69     ValueError: 00 01 02 03
  70     >>> stderr = sys.stderr
  71     >>> sys.stderr = sys.stdout
  72     >>> assert_null(b'\x00\x01\x02\x03', strict=False)
  73     warning: post-data padding not zero: 00 01 02 03
  74     >>> sys.stderr = stderr
  75     """
  76     if buffer and _ord(max(buffer)) != 0:
  77         hex_string = hex_bytes(buffer, spaces=1)
  78         if strict:
  79             raise ValueError(hex_string)
  80         else:
  81             _sys.stderr.write(
  82                 'warning: post-data padding not zero: {}\n'.format(hex_string))
  83
  84 # From ReadWave.c
  85 def byte_order(needToReorderBytes):
  86     little_endian = _sys.byteorder == 'little'
  87     if needToReorderBytes:
  88         little_endian = not little_endian
  89     if little_endian:
  90         return '<'  # little-endian
  91     return '>'  # big-endian
  92
  93 # From ReadWave.c
  94 def need_to_reorder_bytes(version):
  95     # If the low order byte of the version field of the BinHeader
  96     # structure is zero then the file is from a platform that uses
  97     # different byte-ordering and therefore all data will need to be
  98     # reordered.
  99     return version & 0xFF == 0
 100
 101 # From ReadWave.c
 102 def checksum(buffer, byte_order, oldcksum, numbytes):
 103     x = _numpy.ndarray(
 104         (numbytes/2,), # 2 bytes to a short -- ignore trailing odd byte
 105         dtype=_numpy.dtype(byte_order+'h'),
 106         buffer=buffer)
 107     oldcksum += x.sum()
 108     if oldcksum > 2**31:  # fake the C implementation's int rollover
 109         oldcksum %= 2**32
 110         if oldcksum > 2**31:
 111             oldcksum -= 2**31
 112     return oldcksum & 0xffff
 113
 114 def _bytes(obj, encoding='utf-8'):
 115     """Convert bytes or strings into bytes
 116
 117     >>> _bytes(b'123')
 118     '123'
 119     >>> _bytes('123')
 120     '123'
 121     """
 122     if _sys.version_info >= (3,):
 123         if isinstance(obj, bytes):
 124             return obj
 125         else:
 126             return bytes(obj, encoding)
 127     else:
 128         return bytes(obj)