Fix docstrings

temporian/core/event_set_ops.py

@@ -2503,8 +2503,8 @@ def where(
             indexes: ...
             events:
                 (3 events):
-                    timestamps: [ 5. 9. 9.]
-                    'f': ['hello', 'hello', 'goodbye']
+                    timestamps: [5. 9. 9.]
+                    'f': [b'hello' b'hello' b'goodbye']
             ...
 
             ```
@@ -2515,15 +2515,15 @@ def where(
             >>> a = tp.event_set(timestamps=[5, 9, 10],
             ...                  features={'condition': [True, True, False],
             ...                            'yes': [1, 2, 3],
-            ...                            'no': [-1, -2, -3])})
+            ...                            'no': [-1, -2, -3]})
 
             >>> b = a['condition'].where(a['yes'], a['no'])
             >>> b
             indexes: ...
             events:
                 (3 events):
                     timestamps: [ 5. 9. 10.]
-                    'f': [1, 2, -3]
+                    'condition': [ 1 2 -3]
             ...
 
             ```
@@ -2532,16 +2532,16 @@ def where(
 
             ```python
             >>> a = tp.event_set(timestamps=[5, 6, 7, 8, 9],
-            ...                  features={'f': [1, 2, -3, -4, 5],
+            ...                  features={'f': [1, 2, -3, -4, 5]})
 
             >>> # Set values < 0 to nan (cast to float to support nan)
             >>> b = (a['f'] >= 0).where(a['f'].cast(float), np.nan)
             >>> b
             indexes: ...
             events:
                 (5 events):
-                    timestamps: [ 5. 6. 7. 8. 9.]
-                    'f': [ 1., 2., nan, nan, 5.]
+                    timestamps: [5. 6. 7. 8. 9.]
+                    'f': [ 1. 2. nan nan 5.]
             ...
 
             ```

temporian/core/operators/unique_timestamps.py

@@ -14,7 +14,6 @@
 
 
 """Unique timestamps operator class and public API function definitions."""
-
 from temporian.core import operator_lib
 from temporian.core.compilation import compile
 from temporian.core.data.node import (
@@ -23,6 +22,7 @@
 )
 from temporian.core.operators.base import Operator
 from temporian.core.typing import EventSetOrNode
+from temporian.utils.typecheck import typecheck
 from temporian.proto import core_pb2 as pb
 
 
@@ -56,6 +56,7 @@ def build_op_definition(cls) -> pb.OperatorDef:
 operator_lib.register_operator(UniqueTimestamps)
 
 
+@typecheck
 @compile
 def unique_timestamps(input: EventSetOrNode) -> EventSetOrNode:
     assert isinstance(input, EventSetNode)

temporian/core/operators/where.py

@@ -14,7 +14,7 @@
 
 
 """Where operator class and public API function definitions."""
-from typing import Union, Any, Tuple
+from typing import Union, Any
 
 from temporian.core import operator_lib
 from temporian.core.compilation import compile

temporian/implementation/numpy/operators/BUILD

@@ -35,6 +35,7 @@ py_library(
         ":timestamps",
         ":unary",
         ":unique_timestamps",
+        ":where",
         "//temporian/implementation/numpy/operators/binary:arithmetic",
         "//temporian/implementation/numpy/operators/binary:logical",
         "//temporian/implementation/numpy/operators/binary:relational",
@@ -396,5 +397,3 @@ py_library(
         "//temporian/implementation/numpy/data:event_set",
     ],
 )
-
-

temporian/implementation/numpy/operators/__init__.py

@@ -64,3 +64,4 @@
 from temporian.implementation.numpy.operators import tick
 from temporian.implementation.numpy.operators import timestamps
 from temporian.implementation.numpy.operators import unique_timestamps
+from temporian.implementation.numpy.operators import where