Update to use IOTensor for default.

Signed-off-by: Yong Tang <[email protected]>

Author: yongtang

tensorflow_io/core/python/ops/io_tensor.py

@@ -17,11 +17,85 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
+
 import tensorflow as tf
 from tensorflow_io.core.python.ops import core_ops
 from tensorflow_io.kafka.python.ops.kafka_ops import kafka_ops
 
-class IOTensor(object):
+class _IOBaseTensor(object):
+  """_IOBaseTensor"""
+
+  def __init__(self,
+               dtype,
+               shape,
+               properties,
+               internal=False):
+    if not internal:
+      raise ValueError("IOTensor constructor is private; please use one "
+                       "of the factory methods instead (e.g., "
+                       "IOTensor.from_tensor())")
+    self._dtype = dtype
+    self._shape = shape
+    self._properties = collections.OrderedDict(
+        {} if properties is None else properties)
+    super(_IOBaseTensor, self).__init__()
+
+  #=============================================================================
+  # Accessors
+  #=============================================================================
+
+  @property
+  def dtype(self):
+    """The `DType` of values in this tensor."""
+    return self._dtype
+
+  @property
+  def shape(self):
+    """The statically known shape of this io tensor.
+
+    Returns:
+      A `TensorShape` containing the statically known shape of this io
+      tensor. The first dimension could have a size of `None` if this
+      io tensor is from an iterable.
+
+    Examples:
+
+      ```python
+      ```
+    """
+    return self._shape
+
+  @property
+  def rank(self):
+    """The number of dimensions in this io tensor.
+
+    Returns:
+      A Python `int` indicating the number of dimensions in this io
+      tensor.
+    """
+    return tf.rank(self._shape)
+
+  @property
+  def properties(self):
+    """The properties associated with this tensor.
+
+    Returns:
+      A ordered dict with name and properties associated with this tensor.
+    """
+    return self._properties
+
+  #=============================================================================
+  # String Encoding
+  #=============================================================================
+  def __repr__(self):
+    props = "".join([
+        ", %s: %s" % (k, repr(v)) for (k, v) in self.properties.items()])
+    return "<%s: shape=%s, dtype=%s%s>" % (
+        type(self).__name__, self.shape, self.dtype.name, props)
+
+
+class IOTensor(_IOBaseTensor):
   """IOTensor"""
 
   #=============================================================================
@@ -30,6 +104,9 @@ class IOTensor(object):
   def __init__(self,
                dtype,
                shape,
+               resource,
+               function,
+               properties,
                internal=False):
     """Creates an `IOTensor`.
 
@@ -42,19 +119,44 @@ def __init__(self,
     Args:
       dtype: The type of the tensor.
       shape: The shape of the tensor.
+      resource: The resource associated with the IO.
+      function: The function for indexing and accessing items with resource.
+      properties: An ordered dict of properties to be printed.
       internal: True if the constructor is being called by one of the factory
         methods.  If false, an exception will be raised.
 
     Raises:
       TypeError:
     """
-    if not internal:
-      raise ValueError("IOTensor constructor is private; please use one "
-                       "of the factory methods instead (e.g., "
-                       "IOTensor.from_tensor())")
-    self._dtype = dtype
-    self._shape = shape
-    super(IOTensor, self).__init__()
+    self._resource = resource
+    self._function = function
+    super(IOTensor, self).__init__(dtype, shape, properties, internal=internal)
+
+  #=============================================================================
+  # Indexing & Slicing
+  #=============================================================================
+  def __getitem__(self, key):
+    """Returns the specified piece of this IOTensor."""
+    if isinstance(key, slice):
+      start = key.start
+      stop = key.stop
+      step = key.step
+      if start is None:
+        start = 0
+      if stop is None:
+        stop = -1
+      if step is None:
+        step = 1
+    else:
+      start = key
+      stop = key + 1
+      step = 1
+    return self._function(
+        self._resource, start, stop, step, dtype=self.dtype)
+
+  def __len__(self):
+    """Returns the total number of items of this IOTensor."""
+    return abs(self.shape[0])
 
   #=============================================================================
   # Factory Methods
@@ -85,7 +187,7 @@ def from_kafka(cls,
                  subscription,
                  servers=None,
                  timeout=None,
-                 eof=False,
+                 eof=True,
                  conf=None,
                  **kwargs):
     """Creates an `IOTensor` from a Kafka stream.
@@ -114,42 +216,6 @@ def from_kafka(cls,
       return KafkaIOTensor(
           subscription, servers, timeout, eof, conf, internal=True)
 
-  #=============================================================================
-  # Accessors
-  #=============================================================================
-
-  @property
-  def dtype(self):
-    """The `DType` of values in this tensor."""
-    return self._dtype
-
-  @property
-  def shape(self):
-    """The statically known shape of this io tensor.
-
-    Returns:
-      A `TensorShape` containing the statically known shape of this io
-      tensor. The first dimension could have a size of `None` if this
-      io tensor is from an iterable.
-
-    Examples:
-
-      ```python
-      ```
-    """
-    return self._shape
-
-  @property
-  def rank(self):
-    """The number of dimensions in this io tensor.
-
-    Returns:
-      A Python `int` indicating the number of dimensions in this io
-      tensor.
-    """
-    return tf.rank(self._shape)
-
-
   #=============================================================================
   # Tensor Type Conversions
   #=============================================================================
@@ -195,14 +261,7 @@ def to_tensor(self, **kwargs):
     with tf.name_scope(kwargs.get("name", "IOToTensor")):
       return self.__getitem__(slice(None, None))
 
-  #=============================================================================
-  # String Encoding
-  #=============================================================================
-  def __repr__(self):
-    return "<%s: shape=%s, dtype=%s>" % (
-        type(self).__name__, self.shape, self.dtype.name)
-
-class IOIterableTensor(IOTensor):
+class IOIterableTensor(_IOBaseTensor):
   """IOIterableTensor"""
 
   #=============================================================================
@@ -211,62 +270,21 @@ class IOIterableTensor(IOTensor):
   def __init__(self,
                dtype,
                shape,
-               iter_func,
+               function,
+               properties,
                internal=False):
     """Creates an `IOIterableTensor`. """
-    self._iter_func = iter_func
-    super(IOIterableTensor, self).__init__(dtype, shape, internal=internal)
+    self._function = function
+    super(IOIterableTensor, self).__init__(
+        dtype, shape, properties, internal=internal)
 
   #=============================================================================
   # Iterator
   #=============================================================================
   def __iter__(self):
-    return self._iter_func()
-
-class IOIndexableTensor(IOTensor):
-  """IOIndexableTensor"""
-
-  #=============================================================================
-  # Constructor (private)
-  #=============================================================================
-  def __init__(self,
-               dtype,
-               shape,
-               resource,
-               getitem_func,
-               internal=False):
-    """Creates an `IOIndexableTensor`. """
-    self._resource = resource
-    self._getitem_func = getitem_func
-    super(IOIndexableTensor, self).__init__(dtype, shape, internal=internal)
+    return self._function()
 
-  #=============================================================================
-  # Indexing & Slicing
-  #=============================================================================
-  def __getitem__(self, key):
-    """Returns the specified piece of this IOTensor."""
-    if isinstance(key, slice):
-      start = key.start
-      stop = key.stop
-      step = key.step
-      if start is None:
-        start = 0
-      if stop is None:
-        stop = -1
-      if step is None:
-        step = 1
-    else:
-      start = key
-      stop = key + 1
-      step = 1
-    return self._getitem_func(
-        self._resource, start, stop, step, dtype=self.dtype)
-
-  def __len__(self):
-    """Returns the total number of items of this IOTensor."""
-    return abs(self.shape[0])
-
-class AudioIOTensor(IOIndexableTensor):
+class AudioIOTensor(IOTensor):
   """AudioIOTensor"""
 
   #=============================================================================
@@ -281,10 +299,12 @@ def __init__(self,
       dtype = tf.as_dtype(dtypes[0].numpy())
       shape = tf.TensorShape([
           None if dim < 0 else dim for dim in shapes[0].numpy() if dim != 0])
+      properties = collections.OrderedDict({"rate": rate.numpy()})
       self._rate = rate.numpy()
       super(AudioIOTensor, self).__init__(
           dtype, shape,
           resource, core_ops.wav_indexable_get_item,
+          properties,
           internal=internal)
 
   #=============================================================================
@@ -296,14 +316,6 @@ def rate(self):
     """The sampel `rate` of the audio stream"""
     return self._rate
 
-  #=============================================================================
-  # String Encoding
-  #=============================================================================
-  def __repr__(self):
-    return "<%s: shape=%s, dtype=%s, rate=%s>" % (
-        type(self).__name__, self.shape, self.dtype.name, self.rate)
-
-
 class KafkaIOTensor(IOIterableTensor):
   """KafkaIOTensor"""
 
@@ -317,7 +329,7 @@ def __init__(self,
       dtype = tf.string
       shape = tf.TensorShape([None])
 
-      def _iter_func():
+      def function():
         """_iter_func"""
         metadata = []
         if servers is not None:
@@ -326,8 +338,9 @@ def _iter_func():
           metadata.append("conf.timeout=%d" % timeout)
         if eof is not None:
           metadata.append("conf.eof=%d" % (1 if eof else 0))
-        for e in conf:
-          metadata.append(e)
+        if conf is not None:
+          for e in conf:
+            metadata.append(e)
         resource, _, _ = kafka_ops.kafka_iterable_init(
             subscription, metadata=metadata,
             container=scope, shared_name=subscription)
@@ -339,4 +352,4 @@ def _iter_func():
           yield value
 
       super(KafkaIOTensor, self).__init__(
-          dtype, shape, _iter_func, internal=internal)
+          dtype, shape, function, properties=None, internal=internal)