Skip to main content

nautilus_persistence/python/
feather.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Python bindings for the Rust `FeatherWriter` as `StreamingFeatherWriter`.
17
18use std::{
19    cell::RefCell,
20    collections::{HashMap, HashSet},
21    rc::Rc,
22};
23
24use nautilus_common::{
25    live::get_runtime,
26    msgbus::typed_handler::ShareableMessageHandler,
27    python::{cache::PyCache, clock::PyClock},
28};
29use nautilus_core::UnixNanos;
30use nautilus_model::{
31    data::{
32        Bar, Data, FundingRateUpdate, IndexPriceUpdate, InstrumentStatus, MarkPriceUpdate,
33        OptionGreeks, OrderBookDelta, OrderBookDepth10, QuoteTick, TradeTick,
34        close::InstrumentClose,
35    },
36    events::{
37        AccountState, OrderAccepted, OrderCancelRejected, OrderCanceled, OrderDenied,
38        OrderEmulated, OrderExpired, OrderFilled, OrderInitialized, OrderModifyRejected,
39        OrderPendingCancel, OrderPendingUpdate, OrderRejected, OrderReleased, OrderSnapshot,
40        OrderSubmitted, OrderTriggered, OrderUpdated, PositionAdjusted, PositionChanged,
41        PositionClosed, PositionOpened, PositionSnapshot,
42    },
43    python::instruments::pyobject_to_instrument_any,
44    reports::{ExecutionMassStatus, FillReport, OrderStatusReport, PositionStatusReport},
45};
46use object_store::ObjectStoreExt;
47use pyo3::{exceptions::PyIOError, prelude::*};
48
49use crate::{
50    backend::feather::{FeatherWriter, RotationConfig},
51    parquet::{ObjectStoreLocationKind, create_object_store_location_from_path},
52};
53
54/// Python binding for the Rust `FeatherWriter`.
55///
56/// This provides a streaming writer of Nautilus objects into feather files with rotation
57/// capabilities, matching the interface of Python's `StreamingFeatherWriter`.
58#[pyclass(
59    name = "StreamingFeatherWriter",
60    module = "nautilus_trader.core.nautilus_pyo3.persistence",
61    unsendable
62)]
63#[pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.persistence")]
64pub struct PyStreamingFeatherWriter {
65    writer: Rc<RefCell<FeatherWriter>>,
66    handler: Option<ShareableMessageHandler>,
67}
68
69#[pymethods]
70#[pyo3_stub_gen::derive::gen_stub_pymethods]
71impl PyStreamingFeatherWriter {
72    /// Creates a new `StreamingFeatherWriter` instance.
73    ///
74    /// # Parameters
75    ///
76    /// - `path`: The path to persist the stream to. Must be a directory.
77    /// - `cache`: The cache for query info (`PyCache`).
78    /// - `clock`: The clock to use for time-related operations (`PyClock`).
79    /// - `fs_protocol`: Optional filesystem protocol (default: "file").
80    /// - `fs_storage_options`: Optional storage options for cloud backends.
81    /// - `include_types`: Optional list of type names to include (e.g., `["quotes", "trades"]`).
82    /// - `rotation_mode`: Rotation mode (0=SIZE, 1=INTERVAL, `2=SCHEDULED_DATES`, `3=NO_ROTATION`).
83    /// - `max_file_size`: Maximum file size in bytes before rotation (for SIZE mode).
84    /// - `rotation_interval_ns`: Rotation interval in nanoseconds (for `INTERVAL/SCHEDULED_DATES` modes).
85    /// - `rotation_time_ns`: Scheduled rotation time in nanoseconds (for `SCHEDULED_DATES` mode).
86    /// - `flush_interval_ms`: Flush interval in milliseconds (default: 1000). Set to 0 to disable auto-flush.
87    /// - `replace`: If existing files at the given path should be replaced (default: False).
88    #[new]
89    #[pyo3(signature = (
90        path,
91        cache,
92        clock,
93        fs_protocol=None,
94        fs_storage_options=None,
95        include_types=None,
96        rotation_mode=3,
97        max_file_size=1024*1024*1024,
98        rotation_interval_ns=None,
99        rotation_time_ns=None,
100        rotation_timezone="UTC",
101        flush_interval_ms=None,
102        replace=false
103    ))]
104    #[expect(clippy::too_many_arguments, clippy::needless_pass_by_value)]
105    pub fn new(
106        path: String,
107        cache: PyCache,
108        clock: PyClock,
109        fs_protocol: Option<&str>,
110        fs_storage_options: Option<HashMap<String, String>>,
111        include_types: Option<Vec<String>>,
112        rotation_mode: u8,
113        max_file_size: u64,
114        rotation_interval_ns: Option<u64>,
115        rotation_time_ns: Option<u64>,
116        rotation_timezone: &str,
117        flush_interval_ms: Option<u64>,
118        replace: bool,
119    ) -> PyResult<Self> {
120        // Create object store from path
121        // Use fs_protocol to construct the full path if it's a cloud protocol
122        let full_path = match fs_protocol {
123            Some(protocol) if protocol != "file" && !path.contains("://") => {
124                format!("{protocol}://{path}")
125            }
126            _ => path,
127        };
128
129        let storage_options = fs_storage_options
130            .map(|map| map.into_iter().collect::<ahash::AHashMap<String, String>>());
131
132        let location = create_object_store_location_from_path(&full_path, storage_options)
133            .map_err(|e| PyIOError::new_err(format!("Failed to create object store: {e}")))?;
134        let is_local_store = matches!(&location.kind, ObjectStoreLocationKind::Local);
135        let object_store = location.object_store;
136        let base_path = location.base_path;
137
138        // Handle replace parameter - delete existing files if requested
139        if replace {
140            let runtime = get_runtime();
141            let store_ref = object_store.clone();
142            let prefix = if base_path.is_empty() {
143                if !is_local_store {
144                    return Err(PyIOError::new_err(
145                        "replace=True for remote streaming paths requires a non-empty prefix",
146                    ));
147                }
148                None
149            } else {
150                Some(object_store::path::Path::from(base_path.clone()))
151            };
152            runtime
153                .block_on(async {
154                    let mut stream = store_ref.list(prefix.as_ref());
155                    let mut to_delete = Vec::new();
156
157                    while let Some(result) = futures::StreamExt::next(&mut stream).await {
158                        if let Ok(meta) = result {
159                            to_delete.push(meta.location);
160                        }
161                    }
162
163                    for path in to_delete {
164                        let _ = store_ref.delete(&path).await;
165                    }
166                    Ok::<(), anyhow::Error>(())
167                })
168                .map_err(|e| {
169                    PyIOError::new_err(format!("Failed to replace existing files: {e}"))
170                })?;
171        }
172
173        // Convert rotation mode to RotationConfig
174        // Python RotationMode: 0=SIZE, 1=INTERVAL, 2=SCHEDULED_DATES, 3=NO_ROTATION
175        let rotation_config = match rotation_mode {
176            0 => RotationConfig::Size {
177                max_size: max_file_size,
178            },
179            1 => {
180                let interval = rotation_interval_ns.unwrap_or(86_400_000_000_000); // Default 1 day
181                RotationConfig::Interval {
182                    interval_ns: interval,
183                }
184            }
185            2 => {
186                let interval = rotation_interval_ns.unwrap_or(86_400_000_000_000); // Default 1 day
187                let tz = rotation_timezone.parse::<chrono_tz::Tz>().map_err(|e| {
188                    PyIOError::new_err(format!("Failed to parse rotation_timezone: {e}"))
189                })?;
190                let time_ns = rotation_time_ns.unwrap_or(0);
191                RotationConfig::ScheduledDates {
192                    interval_ns: interval,
193                    rotation_time: UnixNanos::from(time_ns),
194                    rotation_timezone: tz,
195                }
196            }
197            _ => RotationConfig::NoRotation, // Default to no rotation for invalid values
198        };
199
200        // Convert include_types to HashSet
201        let type_filter = include_types.map(|types| types.into_iter().collect::<HashSet<String>>());
202
203        let mut per_instrument_types = HashSet::new();
204        per_instrument_types.insert("bars".to_string());
205        per_instrument_types.insert("funding_rate_update".to_string());
206        per_instrument_types.insert("index_prices".to_string());
207        per_instrument_types.insert("mark_prices".to_string());
208        per_instrument_types.insert("order_book_deltas".to_string());
209        per_instrument_types.insert("order_book_depths".to_string());
210        per_instrument_types.insert("option_greeks".to_string());
211        per_instrument_types.insert("quotes".to_string());
212        per_instrument_types.insert("trades".to_string());
213
214        // Extract Clock from Python wrapper
215        // PyClock wraps Rc<RefCell<dyn Clock>>, we get the inner Rc
216        let clock_rc = clock.clock_rc();
217        // Note: Cache parameter is kept for API compatibility with Python StreamingFeatherWriter
218        // but is not directly used by FeatherWriter
219        let _cache = cache;
220
221        // Create FeatherWriter
222        let writer = FeatherWriter::new(
223            base_path,
224            object_store,
225            clock_rc,
226            rotation_config,
227            type_filter,
228            Some(per_instrument_types),
229            flush_interval_ms, // Auto-flush interval in milliseconds
230        );
231
232        Ok(Self {
233            writer: Rc::new(RefCell::new(writer)),
234            handler: None,
235        })
236    }
237
238    /// Subscribes to all messages on the message bus (pattern "*").
239    ///
240    /// This matches the behavior of Python's `StreamingFeatherWriter` when subscribed
241    /// via `trader.subscribe("*", writer.write)`.
242    pub fn subscribe(&mut self) -> PyResult<()> {
243        if self.handler.is_some() {
244            // Already subscribed
245            return Ok(());
246        }
247
248        let handler = FeatherWriter::subscribe_to_message_bus(self.writer.clone())
249            .map_err(|e| PyIOError::new_err(format!("Failed to subscribe to message bus: {e}")))?;
250
251        self.handler = Some(handler);
252        Ok(())
253    }
254
255    /// Unsubscribes from the message bus.
256    pub fn unsubscribe(&mut self) -> PyResult<()> {
257        if let Some(handler) = self.handler.take() {
258            FeatherWriter::unsubscribe_from_message_bus(&handler);
259        }
260        Ok(())
261    }
262
263    /// Writes a data object to the stream.
264    ///
265    /// # Parameters
266    ///
267    /// - `data`: The data object to write (must be a Nautilus data type from pyo3).
268    ///
269    #[expect(
270        clippy::needless_pass_by_value,
271        clippy::too_many_lines,
272        reason = "PyO3 writer binding must downcast supported data variants inline"
273    )]
274    pub fn write(&self, py: Python, data: Py<PyAny>) -> PyResult<()> {
275        macro_rules! try_write {
276            ($type:ty, $name:literal) => {
277                if let Ok(value) = data.extract::<$type>(py) {
278                    let mut writer = self.writer.borrow_mut();
279                    let runtime = get_runtime();
280                    return runtime
281                        .block_on(async { writer.write(value).await })
282                        .map_err(|e| {
283                            PyIOError::new_err(format!("Failed to write {}: {e}", $name))
284                        });
285                }
286            };
287        }
288
289        // Try to convert from common pyo3 data types
290        if let Ok(quote) = data.extract::<QuoteTick>(py) {
291            let mut writer = self.writer.borrow_mut();
292            let runtime = get_runtime();
293            return runtime
294                .block_on(async { writer.write_data(Data::Quote(quote)).await })
295                .map_err(|e| PyIOError::new_err(format!("Failed to write QuoteTick: {e}")));
296        }
297
298        if let Ok(trade) = data.extract::<TradeTick>(py) {
299            let mut writer = self.writer.borrow_mut();
300            let runtime = get_runtime();
301            return runtime
302                .block_on(async { writer.write_data(Data::Trade(trade)).await })
303                .map_err(|e| PyIOError::new_err(format!("Failed to write TradeTick: {e}")));
304        }
305
306        if let Ok(bar) = data.extract::<Bar>(py) {
307            let mut writer = self.writer.borrow_mut();
308            let runtime = get_runtime();
309            return runtime
310                .block_on(async { writer.write_data(Data::Bar(bar)).await })
311                .map_err(|e| PyIOError::new_err(format!("Failed to write Bar: {e}")));
312        }
313
314        if let Ok(delta) = data.extract::<OrderBookDelta>(py) {
315            let mut writer = self.writer.borrow_mut();
316            let runtime = get_runtime();
317            return runtime
318                .block_on(async { writer.write_data(Data::Delta(delta)).await })
319                .map_err(|e| PyIOError::new_err(format!("Failed to write OrderBookDelta: {e}")));
320        }
321
322        if let Ok(depth) = data.extract::<OrderBookDepth10>(py) {
323            let mut writer = self.writer.borrow_mut();
324            let runtime = get_runtime();
325            return runtime
326                .block_on(async { writer.write_data(Data::Depth10(Box::new(depth))).await })
327                .map_err(|e| PyIOError::new_err(format!("Failed to write OrderBookDepth10: {e}")));
328        }
329
330        if let Ok(price) = data.extract::<IndexPriceUpdate>(py) {
331            let mut writer = self.writer.borrow_mut();
332            let runtime = get_runtime();
333            return runtime
334                .block_on(async { writer.write_data(Data::IndexPriceUpdate(price)).await })
335                .map_err(|e| PyIOError::new_err(format!("Failed to write IndexPriceUpdate: {e}")));
336        }
337
338        if let Ok(price) = data.extract::<MarkPriceUpdate>(py) {
339            let mut writer = self.writer.borrow_mut();
340            let runtime = get_runtime();
341            return runtime
342                .block_on(async { writer.write_data(Data::MarkPriceUpdate(price)).await })
343                .map_err(|e| PyIOError::new_err(format!("Failed to write MarkPriceUpdate: {e}")));
344        }
345
346        if let Ok(greeks) = data.extract::<OptionGreeks>(py) {
347            let mut writer = self.writer.borrow_mut();
348            let runtime = get_runtime();
349            return runtime
350                .block_on(async { writer.write_data(Data::OptionGreeks(greeks)).await })
351                .map_err(|e| PyIOError::new_err(format!("Failed to write OptionGreeks: {e}")));
352        }
353
354        if let Ok(close) = data.extract::<InstrumentClose>(py) {
355            let mut writer = self.writer.borrow_mut();
356            let runtime = get_runtime();
357            return runtime
358                .block_on(async { writer.write_data(Data::InstrumentClose(close)).await })
359                .map_err(|e| PyIOError::new_err(format!("Failed to write InstrumentClose: {e}")));
360        }
361
362        try_write!(FundingRateUpdate, "FundingRateUpdate");
363        try_write!(InstrumentStatus, "InstrumentStatus");
364        try_write!(AccountState, "AccountState");
365        try_write!(OrderInitialized, "OrderInitialized");
366        try_write!(OrderDenied, "OrderDenied");
367        try_write!(OrderEmulated, "OrderEmulated");
368        try_write!(OrderSubmitted, "OrderSubmitted");
369        try_write!(OrderAccepted, "OrderAccepted");
370        try_write!(OrderRejected, "OrderRejected");
371        try_write!(OrderPendingCancel, "OrderPendingCancel");
372        try_write!(OrderCanceled, "OrderCanceled");
373        try_write!(OrderCancelRejected, "OrderCancelRejected");
374        try_write!(OrderExpired, "OrderExpired");
375        try_write!(OrderTriggered, "OrderTriggered");
376        try_write!(OrderPendingUpdate, "OrderPendingUpdate");
377        try_write!(OrderReleased, "OrderReleased");
378        try_write!(OrderModifyRejected, "OrderModifyRejected");
379        try_write!(OrderUpdated, "OrderUpdated");
380        try_write!(OrderFilled, "OrderFilled");
381        try_write!(PositionOpened, "PositionOpened");
382        try_write!(PositionChanged, "PositionChanged");
383        try_write!(PositionClosed, "PositionClosed");
384        try_write!(PositionAdjusted, "PositionAdjusted");
385        try_write!(OrderSnapshot, "OrderSnapshot");
386        try_write!(PositionSnapshot, "PositionSnapshot");
387        try_write!(OrderStatusReport, "OrderStatusReport");
388        try_write!(FillReport, "FillReport");
389        try_write!(PositionStatusReport, "PositionStatusReport");
390        try_write!(ExecutionMassStatus, "ExecutionMassStatus");
391
392        // Try instrument types (uses type_str attribute for dispatch)
393        if let Ok(instrument) = pyobject_to_instrument_any(py, data.clone_ref(py)) {
394            let mut writer = self.writer.borrow_mut();
395            let runtime = get_runtime();
396            return runtime
397                .block_on(async { writer.write_instrument(instrument).await })
398                .map_err(|e| PyIOError::new_err(format!("Failed to write instrument: {e}")));
399        }
400
401        Err(PyIOError::new_err(
402            "Unsupported data type for feather writer",
403        ))
404    }
405
406    /// Flushes all active buffers by writing any remaining buffered bytes to the object store.
407    ///
408    /// This is called automatically based on `flush_interval_ms` if configured, but can also
409    /// be called manually by the client.
410    pub fn flush(&self) -> PyResult<()> {
411        let mut writer = self.writer.borrow_mut();
412        let runtime = get_runtime();
413
414        runtime
415            .block_on(async { writer.flush().await })
416            .map_err(|e| PyIOError::new_err(format!("Failed to flush: {e}")))
417    }
418
419    /// Closes all writers by flushing and removing them.
420    ///
421    /// After calling this, no further writes should be performed.
422    pub fn close(&self) -> PyResult<()> {
423        let mut writer = self.writer.borrow_mut();
424        let runtime = get_runtime();
425
426        runtime
427            .block_on(async { writer.close().await })
428            .map_err(|e| PyIOError::new_err(format!("Failed to close: {e}")))
429    }
430
431    /// Returns whether the writer has been closed (no active writers).
432    #[getter]
433    #[must_use]
434    pub fn is_closed(&self) -> bool {
435        self.writer.borrow().is_closed()
436    }
437
438    /// Returns information about the current files being written.
439    ///
440    /// Returns a dictionary mapping writer keys to (size, path) tuples.
441    #[must_use]
442    pub fn get_current_file_info(&self) -> HashMap<String, (u64, String)> {
443        self.writer.borrow().get_current_file_info()
444    }
445
446    /// Returns the next rotation time for a writer, or None if not set.
447    #[pyo3(signature = (type_str, instrument_id=None))]
448    #[must_use]
449    pub fn get_next_rotation_time(
450        &self,
451        type_str: &str,
452        instrument_id: Option<&str>,
453    ) -> Option<u64> {
454        self.writer
455            .borrow()
456            .get_next_rotation_time(type_str, instrument_id)
457            .map(|ns| ns.as_u64())
458    }
459}