pcap_ops.py

"""PcapDataset"""
import tensorflow as tf
from tensorflow_io.core.python.ops import data_ops as data_ops
from tensorflow_io import _load_library
pcap_ops = _load_library('_pcap_ops.so')


class PcapDataset(data_ops.Dataset):
  """A pcap Dataset. Pcap is a popular file format for capturing network packets.
  """

  def __init__(self, filenames, batch=None):
    """Create a pcap Reader.
    Args:
      filenames: A `tf.string` tensor containing one or more filenames.
    """
    batch = 0 if batch is None else batch
    dtypes = [tf.float64, tf.string]
    shapes = [
        tf.TensorShape([]), tf.TensorShape([])] if batch == 0 else [
            tf.TensorShape([None]), tf.TensorShape([None])]
    super(PcapDataset, self).__init__(
        pcap_ops.pcap_dataset,
        pcap_ops.pcap_input(filenames),
        batch, dtypes, shapes)