Add support for list/dict baskets

Signed-off-by: Arham Chopra <arham.chopra@cubistsystematic.com>

Author: arhamchopra

csp_gateway/server/modules/filedrop/adapter.py

@@ -3,14 +3,15 @@
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum, auto
-from typing import Any, Dict, List, get_args, get_origin
+from typing import Any, Dict, List, Optional, TypeVar, get_args, get_origin
 
 import orjson
 import pyarrow.parquet as pq
 from csp import ts
 from csp.impl.pushadapter import PushInputAdapter
 from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
 from csp.impl.wiring import py_push_adapter_def
+from pydantic import TypeAdapter
 from watchdog.events import FileSystemEvent, FileSystemEventHandler
 from watchdog.observers import Observer
 
@@ -49,21 +50,35 @@ class FileDropAdapterConfiguration:
 class FileReaderBase:
     """The base file reader that reads data from files and generates structs"""
 
-    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: object):
+    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: object, deserializer: Optional[object] = None):
         self.field_map = config.field_map
         self.extensions = config.extensions
-        normalized_type = ContainerTypeNormalizer.normalize_type(ts_typ)
-        self.is_list = get_origin(normalized_type) is list
-        if self.is_list:
-            inner_type = get_args(normalized_type)[0]
-            type_adapter = inner_type.type_adapter()
-        else:
-            type_adapter = ts_typ.type_adapter()
-        self.type_adapter = type_adapter
         if hasattr(config, "type_adapter_args"):
             self.context = config.type_adapter_args
         else:
             self.context = {}
+        if not deserializer:
+            normalized_type = ContainerTypeNormalizer.normalize_type(ts_typ)
+            type_adapter = TypeAdapter(normalized_type)
+            if get_origin(normalized_type) is list:
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    data = [apply_field_map(d) for d in data]
+                    return type_adapter.validate_python(data, context=context)
+            elif get_origin(normalized_type) is dict:
+                key_type, inner_type = get_args(normalized_type)
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    data = {k: apply_field_map(d) for k, d in data.items()}
+                    return type_adapter.validate_python(data, context=context)
+            else:
+
+                def deserialize_tick(data, type_adapter=type_adapter, apply_field_map=self.apply_field_map, context=self.context):
+                    return type_adapter.validate_python(apply_field_map(data), context=context)
+
+            self.deserializer = deserialize_tick
+        else:
+            self.deserializer = deserializer
 
     def read(self, src_path: str) -> object:
         """Generator to return stucts from a filepath"""
@@ -74,12 +89,9 @@ def read(self, src_path: str) -> object:
                 should_read = False
         if should_read:
             dicts = self.read_impl(src_path)
-            structs = [self.deserialize_dict(self.apply_field_map(v)) for v in dicts]
-            if self.is_list:
-                yield structs
-            else:
-                for s in structs:
-                    yield s
+            structs = [self.deserializer(d) for d in dicts]
+            for s in structs:
+                yield s
 
     def read_impl(self, src_path: str) -> List[dict]:
         """File type specific implementation"""
@@ -98,11 +110,6 @@ def apply_field_map(self, data: dict) -> dict:
         else:
             return data
 
-    def deserialize_dict(self, dict: dict) -> object:
-        """Convert a dict to struct"""
-
-        return self.type_adapter.validate_python(dict, context=self.context)
-
 
 class FileReaderCsv(FileReaderBase):
     """File reader for json file type"""
@@ -170,18 +177,22 @@ def on_closed(self, event: FileSystemEvent):
                 log.error(f"Failed to read data from {file_path} with exception: {e}, skipping")
 
 
+T = TypeVar("T")
+
+
 class _FileDropImpl(PushInputAdapter):
     FILEREADER_MAP = {
         FileDropType.CSV: FileReaderCsv,
         FileDropType.JSON: FileReaderJson,
         FileDropType.PARQUET: FileReaderParquet,
     }
 
-    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: "T"):  # noqa
+    def __init__(self, config: FileDropAdapterConfiguration, ts_typ: T, deserializer: Optional[object] = None):
+        # NOTE ts_typ is assumed to be a List["Y"] type where "Y" is the actual type
         self.dir_path = config.dir_path
         self.observer = Observer()
         reader = self.FILEREADER_MAP[config.filedrop_type]
-        file_reader = reader(config, ts_typ)
+        file_reader = reader(config, ts_typ[0], deserializer)
         self.event_handler = EventHandlerCustom(self, file_reader)
         self.observer.schedule(self.event_handler, self.dir_path, recursive=False)
 
@@ -195,10 +206,18 @@ def stop(self):
             self.observer.join()
 
 
+# NOTE: Ref: https://github.com/Point72/csp/issues/569
+# Instead of passing the actual type T we want as ts_typ, we have to pass in List[T]
+# This is due to a type normalization bug in csp, where it converts Dict[K, T] to dict, i.e removing the type annotation hints
+# The type annotations are needed to properly deserialize the raw data using the pydantic type adapters.
+# Further we can only take the object type as output due to issues caused by using List[T] instead of T and bypassing
+# the type normalization process in csp. The user of the adapter, needs to keep track of the type they pass in ts_type,
+# and cast the ts[object] to the required ts[T], they want using csp.appy
 filedrop_adapter_def = py_push_adapter_def(
     "filedrop_adapter_def",
     _FileDropImpl,
-    ts["T"],
+    ts[object],
     config=FileDropAdapterConfiguration,
-    ts_typ="T",
+    ts_typ=List[T],
+    deserializer=Optional[object],
 )

csp_gateway/server/modules/filedrop/filedrop.py

@@ -1,9 +1,11 @@
 import logging
 import os
-from typing import Dict, List
+from typing import Dict, List, TypeVar, get_args, get_origin
 
 import csp
 from ccflow import BaseModel
+from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
+from csp.impl.types.tstype import isTsType
 from pydantic import Field
 
 from csp_gateway.server import (
@@ -40,23 +42,39 @@ class ReadFileDropConfiguration(BaseModel):
     )
 
 
+T = TypeVar("T")
+K = TypeVar("K")
+
+
 class ReadFileDrop(GatewayModule):
     """The module to read data from files dropped in specific directories and send them as structs to specific channels"""
 
     directory_configs: Dict[str, List[ReadFileDropConfiguration]] = Field(
         description="Mapping of directories to a list of ReadFileDropConfiguration, for that directory in the filedrop module"
     )
 
+    @csp.node
+    def handle_list_basket(self, data: csp.ts[List[T]], list_size: int) -> csp.OutputBasket(List[csp.ts[T]], shape="list_size"):
+        if csp.ticked(data):
+            return data
+
+    @csp.node
+    def handle_dict_basket(self, data: csp.ts[Dict[K, T]], dict_keys: List[K]) -> csp.OutputBasket(Dict[K, csp.ts[T]], shape="dict_keys"):
+        if csp.ticked(data):
+            return data
+
     def connect(self, channels: GatewayChannels):
         channel_data = {}
+        channel_basket_types = {}
+        channel_base_types = {}
         for dir, configs in self.directory_configs.items():
             # ensure that directories exist
             os.makedirs(dir, exist_ok=True)
             for config in configs:
                 channel_data[config.channel_name] = []
         for dir, configs in self.directory_configs.items():
             for config in configs:
-                type_adapter_args = {
+                context = {
                     "force_new_id": not config.subscribe_with_struct_id,
                     "force_new_timestamp": not config.subscribe_with_struct_timestamp,
                 }
@@ -65,10 +83,40 @@ def connect(self, channels: GatewayChannels):
                     filedrop_type=config.filedrop_type,
                     field_map=config.field_map,
                     extensions=config.extensions,
-                    type_adapter_args=type_adapter_args,
+                    type_adapter_args=context,
                 )
-                channel_type = channels.get_outer_type(config.channel_name).typ
-                data = filedrop_adapter_def(config=adapter_config, ts_typ=channel_type)
+                channel_type = channels.get_outer_type(config.channel_name)
+                if isTsType(channel_type):
+                    non_ts_type = channel_type.typ
+                    channel_basket_types[config.channel_name] = ""
+                else:
+                    normalized_type = ContainerTypeNormalizer.normalize_type(channel_type)
+                    if get_origin(normalized_type) is list:
+                        inner_type = get_args(normalized_type)[0]
+                        if not isTsType(inner_type):
+                            raise ValueError(f"Channel type for {config.channel_name} should be of the form List[TsType], got: {channel_type}")
+                        non_ts_type = List[inner_type.typ]
+                        channel_basket_types[config.channel_name] = "list"
+                    elif get_origin(normalized_type) is dict:
+                        key_type, inner_type = get_args(normalized_type)
+                        if not isTsType(inner_type):
+                            raise ValueError(
+                                f"Channel type for {config.channel_name} should be of the form Dict[KeyType, TsType], got: {channel_type}"
+                            )
+                        non_ts_type = Dict[key_type, inner_type.typ]
+                        channel_basket_types[config.channel_name] = "dict"
+                    else:
+                        raise Exception(f"Channel type cannot be handled: {channel_type}")
+                channel_base_types[config.channel_name] = non_ts_type
+                # NOTE: We have to pass in the type as [type] because of the type normalization issue in csp: https://github.com/Point72/csp/issues/569
+                data = filedrop_adapter_def(config=adapter_config, ts_typ=[non_ts_type], deserializer=None)
                 channel_data[config.channel_name].append(data)
         for channel_name, data_list in channel_data.items():
-            channels.set_channel(channel_name, csp.flatten(data_list))
+            data = csp.flatten(data_list)
+            typed_data = csp.apply(data, lambda x: x, channel_base_types[channel_name])
+            if channel_basket_types[channel_name] == "list":
+                channels.set_channel(channel_name, self.handle_list_basket(typed_data, list_size=channels.dynamic_keys()[channel_name]))
+            elif channel_basket_types[channel_name] == "dict":
+                channels.set_channel(channel_name, self.handle_dict_basket(typed_data, dict_keys=channels.dynamic_keys()[channel_name]))
+            else:
+                channels.set_channel(channel_name, typed_data)