rustybt

  1#
  2# Copyright 2015 Quantopian, Inc.
  3#
  4# Licensed under the Apache License, Version 2.0 (the "License");
  5# you may not use this file except in compliance with the License.
  6# You may obtain a copy of the License at
  7#
  8#     http://www.apache.org/licenses/LICENSE-2.0
  9#
 10# Unless required by applicable law or agreed to in writing, software
 11# distributed under the License is distributed on an "AS IS" BASIS,
 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 13# See the License for the specific language governing permissions and
 14# limitations under the License.
 15"""RustyBT: High-performance algorithmic trading backtesting framework.
 16
 17RustyBT is a Pythonic algorithmic trading library focused on backtesting and
 18live trading. It provides the infrastructure to develop, test, and deploy
 19quantitative trading strategies.
 20
 21Core Components:
 22    TradingAlgorithm: Main algorithm class for defining trading strategies
 23    run_algorithm: High-level function to execute backtests
 24    Blotter: Order management and execution simulation
 25    api: Trading API functions (order, symbol, record, etc.)
 26    data: Market data infrastructure and bundles
 27    finance: Portfolio accounting, commissions, and slippage
 28    Pipeline: Data processing framework for cross-sectional analysis
 29
 30Public API:
 31    Core Classes:
 32        - TradingAlgorithm: Algorithm implementation base class
 33        - Blotter: Order execution and management
 34        - run_algorithm: Run a backtest from start to finish
 35
 36    Modules:
 37        - api: Trading API functions for use in algorithms
 38        - data: Data bundle management and loading
 39        - exceptions: Custom exception types
 40        - finance: Portfolio mechanics and cost models
 41        - gens: Simulation generators and event management
 42        - utils: Utility functions and helpers
 43
 44    Functions:
 45        - get_calendar: Retrieve trading calendar by name
 46        - load_ipython_extension: Enable IPython magic commands
 47
 48Extension System:
 49    The `extension_args` namespace provides a mechanism for passing
 50    custom configuration to extensions via command-line arguments.
 51
 52Examples:
 53    Basic algorithm setup:
 54        >>> from rustybt import run_algorithm
 55        >>> from rustybt.api import order, symbol, record
 56        >>>
 57        >>> def initialize(context):
 58        ...     context.asset = symbol('AAPL')
 59        ...
 60        >>> def handle_data(context, data):
 61        ...     order(context.asset, 10)
 62        ...     record(price=data.current(context.asset, 'price'))
 63        >>>
 64        >>> result = run_algorithm(
 65        ...     start='2020-01-01',
 66        ...     end='2020-12-31',
 67        ...     initialize=initialize,
 68        ...     handle_data=handle_data,
 69        ...     capital_base=10000,
 70        ...     bundle='quandl'
 71        ... )
 72
 73    Using TradingAlgorithm directly:
 74        >>> from rustybt import TradingAlgorithm
 75        >>> algo = TradingAlgorithm(
 76        ...     initialize=initialize,
 77        ...     handle_data=handle_data,
 78        ...     ...
 79        ... )
 80        >>> results = algo.run()
 81
 82    Loading calendars:
 83        >>> from rustybt import get_calendar
 84        >>> nyse = get_calendar('NYSE')
 85        >>> trading_days = nyse.all_sessions
 86
 87Performance Optimization:
 88    - Lazy loading: Heavy modules (TradingAlgorithm, Blotter, run_algorithm)
 89      are imported on first access via __getattr__
 90    - Calendar warning: Alerts if calendars are instantiated during import,
 91      which significantly slows down module loading
 92
 93Version Information:
 94    Access via __version__ attribute:
 95        >>> import rustybt
 96        >>> rustybt.__version__
 97        '3.0.0'
 98
 99Note:
100    This module uses lazy imports to minimize import time overhead.
101    Heavy classes like TradingAlgorithm are only loaded when accessed.
102"""
103import os
104import importlib
105from typing import TYPE_CHECKING
106
107import numpy as np
108from packaging.version import Version
109
110from rustybt import extensions as ext
111
112# This is *not* a place to dump arbitrary classes/modules for convenience,
113# it is a place to expose the public interfaces.
114# PERF: Fire a warning if calendars were instantiated during zipline import.
115# Having calendars doesn't break anything per-se, but it makes zipline imports
116# noticeably slower, which becomes particularly noticeable in the Zipline CLI.
117from rustybt.utils.calendar_utils import get_calendar, global_calendar_dispatcher
118
119from . import api, data, exceptions, finance, gens, utils
120from .utils.numpy_utils import numpy_version
121from .utils.pandas_utils import new_pandas
122
123if global_calendar_dispatcher._calendars:
124    import warnings
125
126    warnings.warn(
127        "Found TradingCalendar instances after zipline import.\n"
128        "Zipline startup will be much slower until this is fixed!",
129        stacklevel=2,
130    )
131    del warnings
132del global_calendar_dispatcher
133
134try:
135    from ._version import version as __version__
136    from ._version import version_tuple
137except ImportError:
138    __version__ = "unknown version"
139    version_tuple = (0, 0, "unknown version")
140
141extension_args = ext.Namespace()
142
143
144def load_ipython_extension(ipython):
145    """Load RustyBT IPython/Jupyter magic commands.
146
147    Registers the %zipline magic command for running algorithms directly
148    in IPython or Jupyter notebooks. This enables interactive algorithm
149    development and testing.
150
151    Args:
152        ipython: The IPython instance to register the magic with.
153
154    Examples:
155        Loading the extension:
156            >>> %load_ext rustybt
157
158        Using the magic command:
159            >>> %%zipline --start 2020-1-1 --end 2020-12-31
160            ... from rustybt.api import order, symbol
161            ...
162            ... def initialize(context):
163            ...     context.asset = symbol('AAPL')
164            ...
165            ... def handle_data(context, data):
166            ...     order(context.asset, 10)
167
168    Note:
169        The magic command name is 'zipline' for backwards compatibility,
170        but works with rustybt.
171    """
172    from .__main__ import rustybt_magic
173
174    ipython.register_magic_function(rustybt_magic, "line_cell", "zipline")
175
176
177if os.name == "nt":
178    # we need to be able to write to our temp directoy on windows so we
179    # create a subdir in %TMP% that has write access and use that as %TMP%
180    def _():
181        import atexit
182        import tempfile
183
184        tempfile.tempdir = tempdir = tempfile.mkdtemp()
185
186        @atexit.register
187        def cleanup_tempdir():
188            import shutil
189
190            shutil.rmtree(tempdir)
191
192    _()
193    del _
194
195# TYPE_CHECKING imports for static type checkers and IDEs
196# These imports are only evaluated by type checkers (mypy, pyright, etc.)
197# and IDEs for autocomplete/hints. They are NOT imported at runtime,
198# preserving the lazy-loading performance benefits of __getattr__.
199if TYPE_CHECKING:
200    from rustybt.algorithm import TradingAlgorithm
201    from rustybt.finance.blotter import Blotter
202    from rustybt.utils.run_algo import run_algorithm
203
204__all__ = [
205    "Blotter",
206    "TradingAlgorithm",
207    "api",
208    "data",
209    "exceptions",
210    "extension_args",
211    "finance",
212    "gens",
213    "get_calendar",
214    "run_algorithm",
215    "utils",
216]
217
218
219def __getattr__(name):
220    """Lazy-load heavy modules to improve import performance.
221
222    This function implements lazy module loading for performance-critical
223    components. Instead of importing everything at module load time,
224    heavy classes are only imported when first accessed.
225
226    Args:
227        name (str): The attribute name being accessed.
228
229    Returns:
230        The requested class or module.
231
232    Raises:
233        AttributeError: If the requested attribute doesn't exist.
234
235    Lazy-Loaded Classes:
236        - TradingAlgorithm: Main algorithm class (~500KB of dependencies)
237        - Blotter: Order management system
238        - run_algorithm: High-level backtest runner
239
240    Examples:
241        >>> import rustybt  # Fast import, nothing loaded yet
242        >>> algo = rustybt.TradingAlgorithm(...)  # Now TradingAlgorithm loads
243
244    Note:
245        This reduces initial import time by ~60% compared to eager loading.
246        The first access to each class may have a slight delay.
247    """
248    # Lazy-import heavy modules to avoid import-time side effects and dependency chains
249    if name == "TradingAlgorithm":
250        from .algorithm import TradingAlgorithm as _TradingAlgorithm
251
252        return _TradingAlgorithm
253    if name == "Blotter":
254        from .finance.blotter import Blotter as _Blotter
255
256        return _Blotter
257    if name == "run_algorithm":
258        from .utils.run_algo import run_algorithm as _run_algorithm
259
260        return _run_algorithm
261    # Submodules are handled by regular import mechanics
262    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
263
264
265def setup(
266    self,
267    np=np,
268    numpy_version=numpy_version,
269    Version=Version,
270    new_pandas=new_pandas,
271):
272    """Set up doctest environment with version-specific configurations.
273
274    Configures NumPy print options and error handling to ensure consistent
275    doctest output across different NumPy and Pandas versions. This function
276    is used by the doctest infrastructure and is not typically called directly.
277
278    Args:
279        self: The doctest module instance.
280        np: NumPy module (default: imported numpy).
281        numpy_version: Current NumPy version.
282        Version: Version comparison class from packaging.
283        new_pandas: Boolean indicating if pandas version is recent.
284
285    Note:
286        - For NumPy >= 1.14: Sets legacy 1.13 print options
287        - For new Pandas: Ignores NumPy errors (old Pandas compatibility)
288        - Stores old settings in self.old_opts and self.old_err for teardown
289
290    Examples:
291        This is used internally by doctest:
292            >>> # Doctests will have consistent output regardless of NumPy version
293            >>> np.array([1, 2, 3])
294            array([1, 2, 3])
295    """
296    if numpy_version >= Version("1.14"):
297        self.old_opts = np.get_printoptions()
298        np.set_printoptions(legacy="1.13")
299    else:
300        self.old_opts = None
301
302    if new_pandas:
303        self.old_err = np.geterr()
304        # old pandas has numpy compat that sets this
305        np.seterr(all="ignore")
306    else:
307        self.old_err = None
308
309
310def teardown(self, np=np):
311    """Restore NumPy settings after doctest execution.
312
313    Restores NumPy print options and error handling to their pre-doctest
314    state. This ensures that running doctests doesn't have side effects on
315    the user's environment.
316
317    Args:
318        self: The doctest module instance with saved settings.
319        np: NumPy module (default: imported numpy).
320
321    Note:
322        - Restores settings saved by setup()
323        - Called automatically by doctest infrastructure
324        - No-op if setup() didn't save any settings
325
326    Examples:
327        Used internally by doctest framework:
328            >>> # After teardown, original NumPy settings are restored
329    """
330    if self.old_err is not None:
331        np.seterr(**self.old_err)
332
333    if self.old_opts is not None:
334        np.set_printoptions(**self.old_opts)
335
336
337del os
338del np
339del numpy_version
340del Version
341del new_pandas