From b1f91643643a04a91302b3d18df14311b8abc9ec Mon Sep 17 00:00:00 2001
From: zombie-einstein <13398815+zombie-einstein@users.noreply.github.com>
Date: Wed, 21 Feb 2024 01:10:56 +0000
Subject: [PATCH] Sphinx doctests and more Rust doc examples (#3)

* Use doctest for sphinx examples

* Test docstring examples

* Add env crate examples

* Additional examples

* Switch docs task ordering
---
 .github/workflows/pre_merge.yaml   |  2 +
 crates/macros/src/lib.rs           |  8 +++
 crates/order_book/src/lib.rs       | 26 ++++++++-
 crates/order_book/src/orderbook.rs | 21 ++++++-
 crates/order_book/src/types.rs     |  2 +-
 crates/step_sim/src/agents/mod.rs  | 94 +++++++++++++++++++++++++++++-
 crates/step_sim/src/env.rs         | 30 +++++++++-
 crates/step_sim/src/lib.rs         | 52 +++++++++++++----
 crates/step_sim/src/runner.rs      | 25 +++++++-
 docs/source/conf.py                |  1 +
 docs/source/pages/example.rst      | 30 ++++++----
 docs/source/pages/usage.rst        | 18 +++---
 pyproject.toml                     |  1 +
 rust/src/order_book.rs             |  9 ++-
 rust/src/step_sim.rs               | 17 ++++--
 src/bourse/step_sim.py             | 19 ++++++
 16 files changed, 312 insertions(+), 43 deletions(-)

diff --git a/.github/workflows/pre_merge.yaml b/.github/workflows/pre_merge.yaml
index b04f0c5..601876d 100644
--- a/.github/workflows/pre_merge.yaml
+++ b/.github/workflows/pre_merge.yaml
@@ -95,3 +95,5 @@ jobs:
         run: pip install hatch
       - name: Install bourse and build docs 📚
         run: hatch run docs:build
+      - name: Test doc examples
+        run: hatch run docs:test
diff --git a/crates/macros/src/lib.rs b/crates/macros/src/lib.rs
index e123238..d1b99f9 100644
--- a/crates/macros/src/lib.rs
+++ b/crates/macros/src/lib.rs
@@ -1,6 +1,14 @@
 use proc_macro::TokenStream;
 use quote::quote;
 
+/// Agent iteration macro
+///
+/// Implements the `AgentSet` trait for a struct
+/// with fields of agent types. It's often the case
+/// we want to implement `update` function that
+/// iterates over a heterogeneous set of agents,
+/// which this macro automates.
+///
 #[proc_macro_derive(Agents)]
 pub fn agents_derive(input: TokenStream) -> TokenStream {
     let ast = syn::parse(input).unwrap();
diff --git a/crates/order_book/src/lib.rs b/crates/order_book/src/lib.rs
index 8b86408..2fbe369 100644
--- a/crates/order_book/src/lib.rs
+++ b/crates/order_book/src/lib.rs
@@ -13,14 +13,38 @@
 //! use bourse_book::types;
 //!
 //! let mut book = bourse_book::OrderBook::new(0, true);
+//!
+//! // Create a new order
 //! let order_id = book.create_order(
 //!     types::Side::Bid, 50, 101, Some(50)
 //! );
+//!
+//! // Place the order on the market
 //! book.place_order(order_id);
+//!
+//! // Get the current touch prices
 //! let (bid, ask) = book.bid_ask();
+//!
+//! // Cancel the order
 //! book.cancel_order(order_id);
 //! ```
-
+//! # Notes
+//!
+//! - Orders are sorted by price-time priority. To
+//!   reduce any ambiguity in ordering the simulated
+//!   time of the market should be updated in
+//!   between placing orders on the market.
+//! - For accuracy prices are stored as unsigned
+//!   integers (as opposed to a float type), hence
+//!   prices from data should be scaled based on
+//!   market tick-size
+//! - Simulated orders are intended to be owned by
+//!   the order book, from which agents/users can
+//!   retrieve order data. Creating an order with
+//!   [OrderBook::create_order] initialises a new
+//!   order entry, but does not immediately place
+//!   the order on the market.
+//!
 mod orderbook;
 mod side;
 pub mod types;
diff --git a/crates/order_book/src/orderbook.rs b/crates/order_book/src/orderbook.rs
index 1ea40e8..a395ef2 100644
--- a/crates/order_book/src/orderbook.rs
+++ b/crates/order_book/src/orderbook.rs
@@ -36,6 +36,22 @@ pub struct OrderEntry {
 }
 
 /// Order book with order and trade history
+///
+/// # Examples
+///
+/// ```
+/// use bourse_book;
+/// use bourse_book::types;
+///
+/// let mut book = bourse_book::OrderBook::new(0, true);
+/// let order_id = book.create_order(
+///     types::Side::Bid, 50, 101, Some(50)
+/// );
+/// book.place_order(order_id);
+/// let (bid, ask) = book.bid_ask();
+/// book.cancel_order(order_id);
+/// ```
+///
 pub struct OrderBook {
     /// Simulated time, intended to represent
     /// nano-seconds, but arbitrary units can
@@ -69,8 +85,9 @@ impl OrderBook {
     ///
     /// # Arguments
     ///
-    /// - `start_time` - Time to assign to the order book
-    /// - `trading` - FLag to indicate if trades will be
+    /// - `start_time` - Simulated time to assign to the
+    ///   order book
+    /// - `trading` - Flag to indicate if trades will be
     ///   executed
     pub fn new(start_time: Nanos, trading: bool) -> Self {
         Self {
diff --git a/crates/order_book/src/types.rs b/crates/order_book/src/types.rs
index 2271bbb..95c15a7 100644
--- a/crates/order_book/src/types.rs
+++ b/crates/order_book/src/types.rs
@@ -1,4 +1,4 @@
-//! Type aliases and order data structures
+//! Type aliases and order data-structures
 
 /// Order-id
 pub type OrderId = usize;
diff --git a/crates/step_sim/src/agents/mod.rs b/crates/step_sim/src/agents/mod.rs
index 2f7e366..3251224 100644
--- a/crates/step_sim/src/agents/mod.rs
+++ b/crates/step_sim/src/agents/mod.rs
@@ -12,11 +12,103 @@ pub use bourse_macros::Agents;
 pub use momentum_agent::MomentumAgent;
 pub use noise_agent::NoiseAgent;
 
+/// Homogeneous agent set functionality
+///
+/// A set of agents that implement this trait
+/// can then be included in a stuct using the
+/// [Agents] macro to combine multiple agent
+/// types.
+///
+/// # Examples
+///
+/// ```
+/// use bourse_de::Env;
+/// use bourse_de::agents::{Agent, Agents, AgentSet};
+/// use fastrand::Rng;
+///
+/// struct AgentType{}
+///
+/// impl Agent for AgentType {
+///     fn update(
+///         &mut self, env: &mut Env, _rng: &mut Rng
+///     ) {}
+/// }
+///
+/// #[derive(Agents)]
+/// struct MixedAgents {
+///     a: AgentType, b: AgentType
+/// }
+/// ```
 pub trait Agent {
+    /// Update the state of the agent(s)
+    ///
+    /// # Argument
+    ///
+    /// - `env` - Reference to a [Env] simulation environment
+    /// - `rng` - Fastrand random generator
+    ///
     fn update(&mut self, env: &mut Env, rng: &mut Rng);
 }
 
 /// Functionality required for simulation agents
+///
+/// Simulation agents provided as an argument to
+/// [crate::sim_runner] must implement this trait,
+/// but the details of the implementation are
+/// left to the user.
+///
+/// It's a common case that we want update to update
+/// a heterogeneous set of agents which can be
+/// automatically implemented with the [Agents] macro
+/// as long as the agent types implement the [Agent]
+/// trait.
+///
+/// # Examples
+///
+/// ```
+/// use bourse_de::Env;
+/// use bourse_de::agents::{Agent, Agents, AgentSet};
+/// use fastrand::Rng;
+///
+/// struct AgentType{}
+///
+/// impl Agent for AgentType {
+///     fn update(
+///         &mut self, env: &mut Env, _rng: &mut Rng
+///     ) {}
+/// }
+///
+/// #[derive(Agents)]
+/// struct MixedAgents {
+///     a: AgentType,
+///     b: AgentType
+/// }
+/// ```
+///
+/// this is equivelant to
+///
+/// ```
+/// # use bourse_de::Env;
+/// # use bourse_de::agents::{Agent, Agents, AgentSet};
+/// # use fastrand::Rng;
+/// # struct AgentType{}
+/// # impl Agent for AgentType {
+/// #    fn update(
+/// #        &mut self, env: &mut Env, _rng: &mut Rng
+/// #     ) {}
+/// # }
+/// struct MixedAgents {
+///     a: AgentType,
+///     b: AgentType
+/// }
+///
+/// impl AgentSet for MixedAgents {
+///     fn update(&mut self, env: &mut Env, rng: &mut Rng){
+///         self.a.update(env, rng);
+///         self.b.update(env, rng);
+///     }
+/// }
+/// ```
 pub trait AgentSet {
     /// Update function called each simulated step
     ///
@@ -26,7 +118,7 @@ pub trait AgentSet {
     ///
     /// The implementing struct is flexible in what
     /// it represent, from a single agent to a group
-    /// of multiple agent types
+    /// of multiple agent types.
     ///
     /// # Arguments
     ///
diff --git a/crates/step_sim/src/env.rs b/crates/step_sim/src/env.rs
index 17384fe..b13a48b 100644
--- a/crates/step_sim/src/env.rs
+++ b/crates/step_sim/src/env.rs
@@ -9,7 +9,35 @@ use bourse_book::OrderBook;
 use fastrand::Rng;
 use std::mem;
 
-// Simulation environment
+/// Discrete event simulation environment
+///
+/// Simulation environment designed for use in a
+/// discrete event simulation. Allows agents/users
+/// to submit order instructions, update
+/// the state of the simulation, and record the
+/// market data.
+///
+/// # Examples
+///
+/// ```
+/// use bourse_de;
+/// use bourse_de::types;
+/// use fastrand::Rng;
+///
+/// let mut env = bourse_de::Env::new(0, 1_000, true);
+/// let mut rng = Rng::with_seed(101);
+///
+/// // Submit a new order instruction
+/// let order_id = env.place_order(
+///     types::Side::Ask,
+///     100,
+///     101,
+///     Some(50),
+/// );
+///
+/// // Update the state of the market
+/// env.step(&mut rng)
+/// ```
 pub struct Env {
     /// Time-length of each simulation step
     step_size: Nanos,
diff --git a/crates/step_sim/src/lib.rs b/crates/step_sim/src/lib.rs
index af6ef9d..d6f4fae 100644
--- a/crates/step_sim/src/lib.rs
+++ b/crates/step_sim/src/lib.rs
@@ -14,19 +14,49 @@
 //! # Examples
 //!
 //! ```
-//! use bourse_de;
-//! use bourse_de::types;
+//! use bourse_book::types::{Price, Side, Vol};
+//! use bourse_de::agents::AgentSet;
+//! use bourse_de::{sim_runner, Env};
 //! use fastrand::Rng;
 //!
-//! let mut env = bourse_de::Env::new(0, 1_000, true);
-//! let mut rng = Rng::with_seed(101);
-//! let order_id = env.place_order(
-//!     types::Side::Ask,
-//!     100,
-//!     101,
-//!     Some(50),
-//! );
-//! env.step(&mut rng)
+//! struct Agents {
+//!     pub offset: Price,
+//!     pub vol: Vol,
+//!     pub n_agents: usize,
+//! }
+//!
+//! impl AgentSet for Agents {
+//!     // Agents place an order on a random side
+//!     // a fixed distance above/below the mid
+//!     fn update(
+//!         &mut self, env: &mut Env, rng: &mut Rng
+//!     ) {
+//!         let (bid, ask) = env.get_orderbook().bid_ask();
+//!         let mid = (ask - bid) / 2;
+//!         let mid_price = bid + mid;
+//!         println!("{}", mid_price);
+//!         for _ in (0..self.n_agents) {
+//!             let side = rng.choice([Side::Bid, Side::Ask]).unwrap();
+//!             match side {
+//!                 Side::Ask => {
+//!                     let p = mid_price - self.offset;
+//!                     println!("==> {}", p);
+//!                     env.place_order(side, self.vol, 101, Some(p));
+//!                 }
+//!                 Side::Bid => {
+//!                     let p = mid_price + self.offset;
+//!                     println!("==< {}", p);
+//!                     env.place_order(side, self.vol, 101, Some(p));
+//!                 }
+//!             }
+//!         }
+//!     }
+//! }
+//!
+//! let mut env = Env::new(0, 1_000_000, true);
+//! let mut agents = Agents{offset: 5, vol: 50, n_agents: 10};
+//!
+//! sim_runner(&mut env, &mut agents, 101, 50);
 //! ```
 
 pub mod agents;
diff --git a/crates/step_sim/src/runner.rs b/crates/step_sim/src/runner.rs
index b55860c..85fd07c 100644
--- a/crates/step_sim/src/runner.rs
+++ b/crates/step_sim/src/runner.rs
@@ -1,4 +1,4 @@
-/// Simulation execution functionality
+//! Simulation execution functionality
 use super::agents::AgentSet;
 use super::env::Env;
 use kdam::tqdm;
@@ -9,6 +9,29 @@ use kdam::tqdm;
 /// in turn can submit instructions to the environment
 /// and then update the environment state)
 ///
+/// # Examples
+///
+/// ```
+/// use bourse_de::{Env, sim_runner};
+/// use bourse_de::agents::AgentSet;
+/// use fastrand::Rng;
+///
+/// // Dummy agent-type
+/// struct Agents{}
+///
+/// impl AgentSet for Agents {
+///     fn update(
+///         &mut self, env: &mut Env, _rng: &mut Rng
+///     ) {}
+/// }
+///
+/// let mut env = bourse_de::Env::new(0, 1_000, true);
+/// let mut agents = Agents{};
+///
+/// // Run for 100 steps from seed 101
+/// sim_runner(&mut env, &mut agents, 101, 100)
+/// ```
+///
 /// # Arguments
 ///
 /// - `env` - Simulation environment
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 111f041..f8f0190 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -11,6 +11,7 @@
     "sphinx_immaterial",
     "sphinx_immaterial.apidoc.python.apigen",
     "sphinx.ext.intersphinx",
+    "sphinx.ext.doctest",
 ]
 
 napoleon_google_docstring = False
diff --git a/docs/source/pages/example.rst b/docs/source/pages/example.rst
index ea46936..626dedf 100644
--- a/docs/source/pages/example.rst
+++ b/docs/source/pages/example.rst
@@ -4,9 +4,16 @@ Simulation Example
 Here we will demonstrate a simple simulation
 where agents randomly place orders.
 
-We will first define the agent class
+First we import Bourse
 
-.. code-block:: python
+.. testcode:: random_example
+
+   import bourse
+
+
+then define an agent class
+
+.. testcode:: random_example
 
    class RandomAgent:
        def __init__(self, i, price_range):
@@ -18,33 +25,36 @@ We will first define the agent class
            # Place an order if not one live
            if self.order_id is None:
                price = rng.integers(*self.price_range)
-               side = bool(rng.random.choice([True, False]))
+               side = bool(rng.choice([True, False]))
                env.place_order(side, 10, self.i, price=price)
            # Cancel live order
            else:
                env.cancel_order(self.order_id)
                self.order_id = None
 
+For an agent to be part of a simulation it should
+define an ``update`` method that takes a
+:py:class:`numpy.random.Generator` and
+:py:class:`bourse.core.StepEnv` as arguments.
+
 In this example an agent places a randomly order if it
 does not have an existing one, and otherwise attempts to
-cancel it's current order. Simulation agents need to define
-an ``update`` function that takes a reference to a
-Numpy random generator and a simulation environment.
+cancel its current order.
 
 We then initialise an environment and set of agents
 
-.. code-block:: python
+.. testcode:: random_example
 
    seed = 101
-   n_steps = 200
+   n_steps = 50
 
-   agents = [RandomAgent(i, (10, 100)) for i in range(100)]
+   agents = [RandomAgent(i, (10, 100)) for i in range(50)]
    env = bourse.core.StepEnv(seed, 0, 100_000)
 
 We can then use :py:meth:`bourse.step_sim.run` to run the
 simulation
 
-.. code-block:: python
+.. testcode:: random_example
 
    market_data = bourse.step_sim.run(env, agents, n_steps, seed)
 
diff --git a/docs/source/pages/usage.rst b/docs/source/pages/usage.rst
index bfbf7ab..5a69d9c 100644
--- a/docs/source/pages/usage.rst
+++ b/docs/source/pages/usage.rst
@@ -18,7 +18,7 @@ Orderbook
 An orderbook is initialised with a start time
 (this is the time used to record events)
 
-.. code-block:: python
+.. testcode:: book_usage
 
    import bourse
 
@@ -27,7 +27,7 @@ An orderbook is initialised with a start time
 The state of the orderbook an then be directly
 updated, for example placing a limit bid order
 
-.. code-block:: python
+.. testcode:: book_usage
 
    order_vol = 10
    trader_id = 101
@@ -37,7 +37,7 @@ updated, for example placing a limit bid order
 
 or cancelling the same order
 
-.. code-block:: python
+.. testcode:: book_usage
 
    book.cancel_order(order_id)
 
@@ -49,7 +49,7 @@ The orderbook also tracks updates, for example
 trades executed on the order book can be
 retrieved with
 
-.. code-block:: python
+.. testcode:: book_usage
 
    trades = book.get_trades()
    # Convert trade data to a dataframe
@@ -58,7 +58,7 @@ retrieved with
    )
 
 See :py:class:`bourse.core.OrderBook`
-for details of the orderbook API.
+for details of the full order book API.
 
 Simulation Environment
 ----------------------
@@ -67,7 +67,9 @@ A simulation environment can be initialised from
 a random seed, start-time, and step-size (i.e. how
 long in time each simulated step is)
 
-.. code-block:: python
+.. testcode:: sim_usage
+
+   import bourse
 
    seed = 101
    step_size = 100_000
@@ -78,7 +80,7 @@ steps, with transactions submitted to a queue to
 be processed at the end of the step. For example
 placing new orders
 
-.. code-block:: python
+.. testcode:: sim_usage
 
    order_id_a = env.place_order(False, 100, 101, price=60)
    order_id_b = env.place_order(True, 100, 101, price=70)
@@ -91,7 +93,7 @@ time to correctly order transactions.
 The simulation environment also tracks market data for each
 step, for example bid-ask prices can be retrieved using
 
-.. code-block:: python
+.. testcode:: sim_usage
 
    bid_prices, ask_prices = env.get_prices()
 
diff --git a/pyproject.toml b/pyproject.toml
index 0b3d3b3..1197941 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -74,3 +74,4 @@ dependencies = [
 
 [tool.hatch.envs.docs.scripts]
 build = "maturin develop && sphinx-build -W -b html docs/source docs/build"
+test = "maturin develop && sphinx-build -W -b doctest docs/source docs/build"
diff --git a/rust/src/order_book.rs b/rust/src/order_book.rs
index a0b40f5..99cab81 100644
--- a/rust/src/order_book.rs
+++ b/rust/src/order_book.rs
@@ -14,17 +14,20 @@ use pyo3::prelude::*;
 /// Examples
 /// --------
 ///
-/// .. code-block:: python
+/// .. testcode:: book_docstring
 ///
 ///    import bourse
 ///
 ///    book = bourse.core.OrderBook(0, True)
+///
 ///    # Place a new order
 ///    order_id = book.place_order(
 ///        True, 100, 0, price=50
 ///    )
+///
 ///    # Get touch prices
 ///    bid, ask = book.bid_ask()
+///
 ///    # Get the status of the order
 ///    status = book.order_status(order_id)
 ///
@@ -72,8 +75,8 @@ impl OrderBook {
     /// When disabled orders can be placed and modified
     /// but will not be matched.
     ///
-    /// Notes
-    /// -----
+    /// Warnings
+    /// --------
     /// There is currently no market uncrossing algorithm
     /// implemented.
     ///
diff --git a/rust/src/step_sim.rs b/rust/src/step_sim.rs
index acc59d3..66acaa8 100644
--- a/rust/src/step_sim.rs
+++ b/rust/src/step_sim.rs
@@ -24,7 +24,7 @@ use pyo3::prelude::*;
 /// Examples
 /// --------
 ///
-/// .. code-block:: python
+/// .. testcode:: step_sim_docstring
 ///
 ///    import bourse
 ///
@@ -36,9 +36,18 @@ use pyo3::prelude::*;
 ///        seed, start_time, step_size
 ///    )
 ///
-///    order_id = env.place_order(True, 100, 99, price=50)
+///    # Create an order to be placed in the
+///    # next update
+///    order_id = env.place_order(
+///        True, 100, 99, price=50
+///    )
+///
+///    # Update the environment
 ///    env.step()
 ///
+///    # Get price history data
+///    bid_price, ask_prices = env.get_prices()
+///
 #[pyclass]
 pub struct StepEnv {
     env: BaseEnv,
@@ -110,8 +119,8 @@ impl StepEnv {
     /// When disabled orders can be placed and modified
     /// but will not be matched.
     ///
-    /// Notes
-    /// -----
+    /// Warnings
+    /// --------
     /// There is currently no market uncrossing algorithm
     /// implemented.
     ///
diff --git a/src/bourse/step_sim.py b/src/bourse/step_sim.py
index 11ae27b..d8753f7 100644
--- a/src/bourse/step_sim.py
+++ b/src/bourse/step_sim.py
@@ -24,6 +24,25 @@ def run(
     are randomly shuffled and process, updating the state of
     the market.
 
+    Examples
+    --------
+
+    .. testsetup:: runner_docstring
+
+       import bourse
+
+       agents = []
+       env = env = bourse.core.StepEnv(0, 0, 1000)
+
+    .. testcode:: runner_docstring
+
+       market_data = bourse.step_sim.run(
+           env,     # Simulation environment
+           agents,  # List of agents
+           50,      # Number of steps
+           101      # Random seed
+       )
+
     Parameters
     ----------
     env: bourse.core.StepEnv