advent_16_part1.py

TAG_LITERAL = 4
TYPE_BITS = 0

class Stream:
  def __init__(self, data):
    self.data = data
    self.ready_carry = 0
    self.carry_size = 0
    self.carry = 0
    self.nibbles = self._init_read()

  def get(self, size):
    data = self.carry
    data_size = self.carry_size
    while data_size < size:
      data = data << 4
      data = data | next(self.nibbles)
      data_size += 4
    if data_size > size:
      self.carry_size = data_size - size
      self.carry = data & (2 ** self.carry_size - 1)
      data = data >> self.carry_size
    else:
      self.carry = 0
      self.carry_size = 0
    return data

  def _init_read(self):
    for char in self.data:
      yield int(char, 16)


def get_literal_value(stream):
  value = 0
  bits_read = 6 # Tag + Version
  while True:
    has_more_data = stream.get(1)
    value = value << 4
    value = value | stream.get(4)
    bits_read += 5
    if not has_more_data:
      break
  return value, bits_read

def decode(input_data):
  stream = Stream(input_data)
  return decode_packet(stream)

def decode_packet(stream):
  version = stream.get(3)
  tag = stream.get(3)
  value = 0
  if tag == TAG_LITERAL:
    value, bits_read = get_literal_value(stream)
    return (bits_read, version, tag, value)
  length_type = stream.get(1)
  subpackages = []
  if length_type == TYPE_BITS:
    bits_length = stream.get(15)
    bits_read = 0
    while bits_read < bits_length:
      subpackages.append(decode_packet(stream))
      bits_read += subpackages[-1][0]
    bits_read += 22 # 3 for version, 3 for tag, 1 length type + 15 for length
    return (bits_read, version, tag, subpackages)
  else: # TYPE_SUBPACKAGES
    package_length = stream.get(11)
    subpackages = []
    bits_read = 0
    while len(subpackages) < package_length:
      subpackages.append(decode_packet(stream))
      bits_read += subpackages[-1][0]
    bits_read += 18 # 3 for version, 3 for tag, 1 length type + 11 package count
    return (bits_read, version, tag, subpackages)

def sum_versions(package):
  length, version, tag, other = package
  total = version
  if tag != TAG_LITERAL:
    for child in other:
      total += sum_versions(child)
  return total

if __name__ == '__main__':
  with open('input.txt') as data:
    package = decode(data.readline().strip())
    print(package, sum_versions(package))