Improve type stub gen for nautilus_trader v2

Author: cjdsellers

.pre-commit-config.yaml

@@ -160,7 +160,7 @@ repos:
     rev: v1.18.2
     hooks:
       - id: mypy
-        exclude: ^python/nautilus_trader/.*__init__\.py$
+        exclude: (^python/|^nautilus_trader/.*\.pyi$)
         args: [
           "--config", "pyproject.toml",
           "--allow-incomplete-defs",

Cargo.lock

@@ -37,6 +37,7 @@ python = [
   "nautilus-network/python",
   "nautilus-system/python",
   "pyo3",
+  "pyo3-stub-gen",
 ]
 
 [package.metadata.docs.rs]
@@ -78,6 +79,7 @@ ustr = { workspace = true }
 hypersync-client = { workspace = true, optional = true }
 hypersync-schema = { workspace = true, optional = true }
 pyo3 = { workspace = true, optional = true }
+pyo3-stub-gen = { workspace = true, optional = true }
 
 [dev-dependencies]
 rstest = { workspace = true }

crates/adapters/blockchain/Cargo.toml

@@ -22,6 +22,10 @@ use nautilus_model::defi::{DexType, SharedChain};
     feature = "python",
     pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.blockchain")
 )]
+#[cfg_attr(
+    feature = "python",
+    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.blockchain")
+)]
 pub struct DexPoolFilters {
     /// Whether to exclude pools containing tokens with empty name or symbol fields.
     pub remove_pools_with_empty_erc20fields: bool,
@@ -52,6 +56,10 @@ impl Default for DexPoolFilters {
     feature = "python",
     pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.blockchain")
 )]
+#[cfg_attr(
+    feature = "python",
+    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.blockchain")
+)]
 pub struct BlockchainDataClientConfig {
     /// The blockchain chain configuration.
     pub chain: SharedChain,

crates/adapters/blockchain/src/config.rs

@@ -38,6 +38,10 @@ impl ClientConfig for BlockchainDataClientConfig {
     feature = "python",
     pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.blockchain")
 )]
+#[cfg_attr(
+    feature = "python",
+    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.adapters.blockchain")
+)]
 pub struct BlockchainDataClientFactory;
 
 impl BlockchainDataClientFactory {

crates/adapters/blockchain/src/factories.rs

@@ -24,6 +24,7 @@ use pyo3::prelude::*;
 use crate::config::{BlockchainDataClientConfig, DexPoolFilters};
 
 #[pymethods]
+#[pyo3_stub_gen::derive::gen_stub_pymethods(module = "nautilus_trader.adapters.blockchain")]
 impl DexPoolFilters {
     /// Creates a new `DexPoolFilters` instance.
     #[new]
@@ -34,13 +35,26 @@ impl DexPoolFilters {
 }
 
 #[pymethods]
+#[pyo3_stub_gen::derive::gen_stub_pymethods(module = "nautilus_trader.adapters.blockchain")]
 impl BlockchainDataClientConfig {
     /// Creates a new `BlockchainDataClientConfig` instance.
     #[new]
     #[allow(clippy::too_many_arguments)]
     #[pyo3(signature = (chain, dex_ids, http_rpc_url, rpc_requests_per_second=None, multicall_calls_per_rpc_request=None, wss_rpc_url=None, use_hypersync_for_live_data=true, from_block=None, pool_filters=None, postgres_cache_database_config=None))]
     fn py_new(
+        #[gen_stub(
+            override_type(
+                type_repr = "nautilus_trader.model.Chain",
+                imports = ("nautilus_trader.model",),
+            ),
+        )]
         chain: &Chain,
+        #[gen_stub(
+            override_type(
+                type_repr = "typing.Sequence[nautilus_trader.model.DexType]",
+                imports = ("typing", "nautilus_trader.model"),
+            ),
+        )]
         dex_ids: Vec<DexType>,
         http_rpc_url: String,
         rpc_requests_per_second: Option<u32>,
@@ -49,6 +63,12 @@ impl BlockchainDataClientConfig {
         use_hypersync_for_live_data: bool,
         from_block: Option<u64>,
         pool_filters: Option<DexPoolFilters>,
+        #[gen_stub(
+            override_type(
+                type_repr = "typing.Optional[nautilus_trader.infrastructure.PostgresConnectOptions]",
+                imports = ("typing", "nautilus_trader.infrastructure"),
+            ),
+        )]
         postgres_cache_database_config: Option<PostgresConnectOptions>,
     ) -> Self {
         Self::new(
@@ -67,6 +87,12 @@ impl BlockchainDataClientConfig {
 
     /// Returns the chain configuration.
     #[getter]
+    #[gen_stub(
+        override_return_type(
+            type_repr = "nautilus_trader.model.Chain",
+            imports = ("nautilus_trader.model",),
+        ),
+    )]
     fn chain(&self) -> Chain {
         (*self.chain).clone()
     }

crates/adapters/blockchain/src/python/config.rs

@@ -42,10 +42,8 @@ use strum::{Display, EnumIter, EnumString};
 )]
 #[non_exhaustive]
 #[strum(ascii_case_insensitive)]
-#[cfg_attr(
-    feature = "python",
-    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.model")
-)]
+#[cfg_attr(feature = "python", pyo3::pyclass(module = "nautilus_trader.model"))]
+#[cfg_attr(feature = "python", pyo3_stub_gen::derive::gen_stub_pyclass_enum)]
 pub enum Blockchain {
     Abstract,
     Arbitrum,
@@ -131,6 +129,7 @@ pub enum Blockchain {
 
 /// Defines a blockchain with its unique identifiers and connection details for network interaction.
 #[cfg_attr(feature = "python", pyo3::pyclass(module = "nautilus_pyo3.model"))]
+#[cfg_attr(feature = "python", pyo3_stub_gen::derive::gen_stub_pyclass)]
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct Chain {
     /// The blockchain network type.

crates/model/src/defi/chain.rs

@@ -38,10 +38,8 @@ use crate::{
     strum::Display,
     strum::EnumIter,
 )]
-#[cfg_attr(
-    feature = "python",
-    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.model")
-)]
+#[cfg_attr(feature = "python", pyo3::pyclass(module = "nautilus_trader.model"))]
+#[cfg_attr(feature = "python", pyo3_stub_gen::derive::gen_stub_pyclass_enum)]
 #[non_exhaustive]
 pub enum AmmType {
     /// Constant Product Automated Market Maker.
@@ -74,10 +72,8 @@ pub enum AmmType {
     Serialize,
     Deserialize,
 )]
-#[cfg_attr(
-    feature = "python",
-    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.model")
-)]
+#[cfg_attr(feature = "python", pyo3::pyclass(module = "nautilus_trader.model"))]
+#[cfg_attr(feature = "python", pyo3_stub_gen::derive::gen_stub_pyclass_enum)]
 pub enum DexType {
     AerodromeSlipstream,
     AerodromeV1,
@@ -107,10 +103,8 @@ impl DexType {
 
 /// Represents a decentralized exchange (DEX) in a blockchain ecosystem.
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
-#[cfg_attr(
-    feature = "python",
-    pyo3::pyclass(module = "nautilus_trader.core.nautilus_pyo3.model")
-)]
+#[cfg_attr(feature = "python", pyo3::pyclass(module = "nautilus_trader.model"))]
+#[cfg_attr(feature = "python", pyo3_stub_gen::derive::gen_stub_pyclass)]
 pub struct Dex {
     /// The blockchain network where this DEX operates.
     pub chain: Chain,

crates/model/src/defi/dex.rs

@@ -261,6 +261,8 @@ no_strict_optional = false
 warn_no_return = true
 warn_unused_configs = true
 warn_unused_ignores = true
+mypy_path = ["."]
+exclude = "(^python/examples/|^examples/)"
 
 [[tool.mypy.overrides]]
 no_strict_optional = true