From 963b2edda5b687e07398e3eb16ed51f55c9f1b57 Mon Sep 17 00:00:00 2001 From: Adithya Balaji Date: Mon, 13 Apr 2020 23:37:54 -0400 Subject: [PATCH] Add Schema loading capabilities to SessionManager get and post * Remove typing-extensions since it is no longer used * move doc --> docs in pyproject.toml * Re-factor SessionManager to leverage get / post method automatic Schema building. * Add beginning of portfolio model. --- poetry.lock | 76 ++++++++++---------- pyproject.toml | 5 +- pyrh/models/__init__.py | 16 ++++- pyrh/models/portfolio.py | 38 ++++++++++ pyrh/models/sessionmanager.py | 131 ++++++++++------------------------ 5 files changed, 130 insertions(+), 136 deletions(-) create mode 100644 pyrh/models/portfolio.py diff --git a/poetry.lock b/poetry.lock index 51715c89..4af021dd 100644 --- a/poetry.lock +++ b/poetry.lock @@ -141,7 +141,7 @@ description = "Code coverage measurement for Python" name = "coverage" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*, <4" -version = "5.0.4" +version = "5.1" [package.extras] toml = ["toml"] @@ -293,7 +293,7 @@ description = "A very fast and expressive template engine." name = "jinja2" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" -version = "2.11.1" +version = "2.11.2" [package.dependencies] MarkupSafe = ">=0.23" @@ -869,7 +869,7 @@ python-versions = "*" version = "1.4.1" [[package]] -category = "main" +category = "dev" description = "Backported and Experimental Type Hints for Python 3.5+" name = "typing-extensions" optional = false @@ -939,10 +939,10 @@ docs = ["sphinx", "jaraco.packaging (>=3.2)", "rst.linker (>=1.9)"] testing = ["jaraco.itertools", "func-timeout"] [extras] -doc = ["sphinx", "sphinx-autodoc-typehints", "sphinx_rtd_theme", "autodocsumm", "tomlkit"] +docs = ["sphinx", "sphinx-autodoc-typehints", "sphinx_rtd_theme", "autodocsumm", "tomlkit"] [metadata] -content-hash = "7c9aae465d2da37302994f38ded870214339a48855fd84023c27b1a32a6f5fe8" +content-hash = "2daf35a72d1823fc01a23dc7e34033b8c227c8e34b06c679abb368560ee9e45d" python-versions = "^3.6" [metadata.files] @@ -998,37 +998,37 @@ colorama = [ {file = "colorama-0.4.3.tar.gz", hash = "sha256:e96da0d330793e2cb9485e9ddfd918d456036c7149416295932478192f4436a1"}, ] coverage = [ - {file = "coverage-5.0.4-cp27-cp27m-macosx_10_12_x86_64.whl", hash = "sha256:8a620767b8209f3446197c0e29ba895d75a1e272a36af0786ec70fe7834e4307"}, - {file = "coverage-5.0.4-cp27-cp27m-macosx_10_13_intel.whl", hash = "sha256:73aa6e86034dad9f00f4bbf5a666a889d17d79db73bc5af04abd6c20a014d9c8"}, - {file = "coverage-5.0.4-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:408ce64078398b2ee2ec08199ea3fcf382828d2f8a19c5a5ba2946fe5ddc6c31"}, - {file = "coverage-5.0.4-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:cda33311cb9fb9323958a69499a667bd728a39a7aa4718d7622597a44c4f1441"}, - {file = "coverage-5.0.4-cp27-cp27m-win32.whl", hash = "sha256:5f587dfd83cb669933186661a351ad6fc7166273bc3e3a1531ec5c783d997aac"}, - {file = "coverage-5.0.4-cp27-cp27m-win_amd64.whl", hash = "sha256:9fad78c13e71546a76c2f8789623eec8e499f8d2d799f4b4547162ce0a4df435"}, - {file = "coverage-5.0.4-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:2e08c32cbede4a29e2a701822291ae2bc9b5220a971bba9d1e7615312efd3037"}, - {file = "coverage-5.0.4-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:922fb9ef2c67c3ab20e22948dcfd783397e4c043a5c5fa5ff5e9df5529074b0a"}, - {file = "coverage-5.0.4-cp35-cp35m-macosx_10_12_x86_64.whl", hash = "sha256:c3fc325ce4cbf902d05a80daa47b645d07e796a80682c1c5800d6ac5045193e5"}, - {file = "coverage-5.0.4-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:046a1a742e66d065d16fb564a26c2a15867f17695e7f3d358d7b1ad8a61bca30"}, - {file = "coverage-5.0.4-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:6ad6ca45e9e92c05295f638e78cd42bfaaf8ee07878c9ed73e93190b26c125f7"}, - {file = "coverage-5.0.4-cp35-cp35m-win32.whl", hash = "sha256:eda55e6e9ea258f5e4add23bcf33dc53b2c319e70806e180aecbff8d90ea24de"}, - {file = "coverage-5.0.4-cp35-cp35m-win_amd64.whl", hash = "sha256:4a8a259bf990044351baf69d3b23e575699dd60b18460c71e81dc565f5819ac1"}, - {file = "coverage-5.0.4-cp36-cp36m-macosx_10_13_x86_64.whl", hash = "sha256:f372cdbb240e09ee855735b9d85e7f50730dcfb6296b74b95a3e5dea0615c4c1"}, - {file = "coverage-5.0.4-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:a37c6233b28e5bc340054cf6170e7090a4e85069513320275a4dc929144dccf0"}, - {file = "coverage-5.0.4-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:443be7602c790960b9514567917af538cac7807a7c0c0727c4d2bbd4014920fd"}, - {file = "coverage-5.0.4-cp36-cp36m-win32.whl", hash = "sha256:165a48268bfb5a77e2d9dbb80de7ea917332a79c7adb747bd005b3a07ff8caf0"}, - {file = "coverage-5.0.4-cp36-cp36m-win_amd64.whl", hash = "sha256:0a907199566269e1cfa304325cc3b45c72ae341fbb3253ddde19fa820ded7a8b"}, - {file = "coverage-5.0.4-cp37-cp37m-macosx_10_13_x86_64.whl", hash = "sha256:513e6526e0082c59a984448f4104c9bf346c2da9961779ede1fc458e8e8a1f78"}, - {file = "coverage-5.0.4-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:3844c3dab800ca8536f75ae89f3cf566848a3eb2af4d9f7b1103b4f4f7a5dad6"}, - {file = "coverage-5.0.4-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:641e329e7f2c01531c45c687efcec8aeca2a78a4ff26d49184dce3d53fc35014"}, - {file = "coverage-5.0.4-cp37-cp37m-win32.whl", hash = "sha256:db1d4e38c9b15be1521722e946ee24f6db95b189d1447fa9ff18dd16ba89f732"}, - {file = "coverage-5.0.4-cp37-cp37m-win_amd64.whl", hash = "sha256:62061e87071497951155cbccee487980524d7abea647a1b2a6eb6b9647df9006"}, - {file = "coverage-5.0.4-cp38-cp38-macosx_10_13_x86_64.whl", hash = "sha256:65a7e00c00472cd0f59ae09d2fb8a8aaae7f4a0cf54b2b74f3138d9f9ceb9cb2"}, - {file = "coverage-5.0.4-cp38-cp38-manylinux1_i686.whl", hash = "sha256:1f66cf263ec77af5b8fe14ef14c5e46e2eb4a795ac495ad7c03adc72ae43fafe"}, - {file = "coverage-5.0.4-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:85596aa5d9aac1bf39fe39d9fa1051b0f00823982a1de5766e35d495b4a36ca9"}, - {file = "coverage-5.0.4-cp38-cp38-win32.whl", hash = "sha256:86a0ea78fd851b313b2e712266f663e13b6bc78c2fb260b079e8b67d970474b1"}, - {file = "coverage-5.0.4-cp38-cp38-win_amd64.whl", hash = "sha256:03f630aba2b9b0d69871c2e8d23a69b7fe94a1e2f5f10df5049c0df99db639a0"}, - {file = "coverage-5.0.4-cp39-cp39-win32.whl", hash = "sha256:7c9762f80a25d8d0e4ab3cb1af5d9dffbddb3ee5d21c43e3474c84bf5ff941f7"}, - {file = "coverage-5.0.4-cp39-cp39-win_amd64.whl", hash = "sha256:4482f69e0701139d0f2c44f3c395d1d1d37abd81bfafbf9b6efbe2542679d892"}, - {file = "coverage-5.0.4.tar.gz", hash = "sha256:1b60a95fc995649464e0cd48cecc8288bac5f4198f21d04b8229dc4097d76823"}, + {file = "coverage-5.1-cp27-cp27m-macosx_10_12_x86_64.whl", hash = "sha256:0cb4be7e784dcdc050fc58ef05b71aa8e89b7e6636b99967fadbdba694cf2b65"}, + {file = "coverage-5.1-cp27-cp27m-macosx_10_13_intel.whl", hash = "sha256:c317eaf5ff46a34305b202e73404f55f7389ef834b8dbf4da09b9b9b37f76dd2"}, + {file = "coverage-5.1-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:b83835506dfc185a319031cf853fa4bb1b3974b1f913f5bb1a0f3d98bdcded04"}, + {file = "coverage-5.1-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:5f2294dbf7875b991c381e3d5af2bcc3494d836affa52b809c91697449d0eda6"}, + {file = "coverage-5.1-cp27-cp27m-win32.whl", hash = "sha256:de807ae933cfb7f0c7d9d981a053772452217df2bf38e7e6267c9cbf9545a796"}, + {file = "coverage-5.1-cp27-cp27m-win_amd64.whl", hash = "sha256:bf9cb9a9fd8891e7efd2d44deb24b86d647394b9705b744ff6f8261e6f29a730"}, + {file = "coverage-5.1-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:acf3763ed01af8410fc36afea23707d4ea58ba7e86a8ee915dfb9ceff9ef69d0"}, + {file = "coverage-5.1-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:dec5202bfe6f672d4511086e125db035a52b00f1648d6407cc8e526912c0353a"}, + {file = "coverage-5.1-cp35-cp35m-macosx_10_12_x86_64.whl", hash = "sha256:7a5bdad4edec57b5fb8dae7d3ee58622d626fd3a0be0dfceda162a7035885ecf"}, + {file = "coverage-5.1-cp35-cp35m-manylinux1_i686.whl", hash = "sha256:1601e480b9b99697a570cea7ef749e88123c04b92d84cedaa01e117436b4a0a9"}, + {file = "coverage-5.1-cp35-cp35m-manylinux1_x86_64.whl", hash = "sha256:dbe8c6ae7534b5b024296464f387d57c13caa942f6d8e6e0346f27e509f0f768"}, + {file = "coverage-5.1-cp35-cp35m-win32.whl", hash = "sha256:a027ef0492ede1e03a8054e3c37b8def89a1e3c471482e9f046906ba4f2aafd2"}, + {file = "coverage-5.1-cp35-cp35m-win_amd64.whl", hash = "sha256:0e61d9803d5851849c24f78227939c701ced6704f337cad0a91e0972c51c1ee7"}, + {file = "coverage-5.1-cp36-cp36m-macosx_10_13_x86_64.whl", hash = "sha256:2d27a3f742c98e5c6b461ee6ef7287400a1956c11421eb574d843d9ec1f772f0"}, + {file = "coverage-5.1-cp36-cp36m-manylinux1_i686.whl", hash = "sha256:66460ab1599d3cf894bb6baee8c684788819b71a5dc1e8fa2ecc152e5d752019"}, + {file = "coverage-5.1-cp36-cp36m-manylinux1_x86_64.whl", hash = "sha256:5c542d1e62eece33c306d66fe0a5c4f7f7b3c08fecc46ead86d7916684b36d6c"}, + {file = "coverage-5.1-cp36-cp36m-win32.whl", hash = "sha256:2742c7515b9eb368718cd091bad1a1b44135cc72468c731302b3d641895b83d1"}, + {file = "coverage-5.1-cp36-cp36m-win_amd64.whl", hash = "sha256:dead2ddede4c7ba6cb3a721870f5141c97dc7d85a079edb4bd8d88c3ad5b20c7"}, + {file = "coverage-5.1-cp37-cp37m-macosx_10_13_x86_64.whl", hash = "sha256:01333e1bd22c59713ba8a79f088b3955946e293114479bbfc2e37d522be03355"}, + {file = "coverage-5.1-cp37-cp37m-manylinux1_i686.whl", hash = "sha256:e1ea316102ea1e1770724db01998d1603ed921c54a86a2efcb03428d5417e489"}, + {file = "coverage-5.1-cp37-cp37m-manylinux1_x86_64.whl", hash = "sha256:adeb4c5b608574a3d647011af36f7586811a2c1197c861aedb548dd2453b41cd"}, + {file = "coverage-5.1-cp37-cp37m-win32.whl", hash = "sha256:782caea581a6e9ff75eccda79287daefd1d2631cc09d642b6ee2d6da21fc0a4e"}, + {file = "coverage-5.1-cp37-cp37m-win_amd64.whl", hash = "sha256:00f1d23f4336efc3b311ed0d807feb45098fc86dee1ca13b3d6768cdab187c8a"}, + {file = "coverage-5.1-cp38-cp38-macosx_10_13_x86_64.whl", hash = "sha256:402e1744733df483b93abbf209283898e9f0d67470707e3c7516d84f48524f55"}, + {file = "coverage-5.1-cp38-cp38-manylinux1_i686.whl", hash = "sha256:a3f3654d5734a3ece152636aad89f58afc9213c6520062db3978239db122f03c"}, + {file = "coverage-5.1-cp38-cp38-manylinux1_x86_64.whl", hash = "sha256:6402bd2fdedabbdb63a316308142597534ea8e1895f4e7d8bf7476c5e8751fef"}, + {file = "coverage-5.1-cp38-cp38-win32.whl", hash = "sha256:8fa0cbc7ecad630e5b0f4f35b0f6ad419246b02bc750de7ac66db92667996d24"}, + {file = "coverage-5.1-cp38-cp38-win_amd64.whl", hash = "sha256:79a3cfd6346ce6c13145731d39db47b7a7b859c0272f02cdb89a3bdcbae233a0"}, + {file = "coverage-5.1-cp39-cp39-win32.whl", hash = "sha256:a82b92b04a23d3c8a581fc049228bafde988abacba397d57ce95fe95e0338ab4"}, + {file = "coverage-5.1-cp39-cp39-win_amd64.whl", hash = "sha256:bb28a7245de68bf29f6fb199545d072d1036a1917dca17a1e75bbb919e14ee8e"}, + {file = "coverage-5.1.tar.gz", hash = "sha256:f90bfc4ad18450c80b024036eaf91e4a246ae287701aaa88eaebebf150868052"}, ] darglint = [ {file = "darglint-1.2.1-py3-none-any.whl", hash = "sha256:16ee69a67fc0f3a89917ba4028b9c50491d7cb4e569cb94eed2e013e2a574c77"}, @@ -1083,8 +1083,8 @@ isort = [ {file = "isort-4.3.21.tar.gz", hash = "sha256:54da7e92468955c4fceacd0c86bd0ec997b0e1ee80d97f67c35a78b719dccab1"}, ] jinja2 = [ - {file = "Jinja2-2.11.1-py2.py3-none-any.whl", hash = "sha256:b0eaf100007721b5c16c1fc1eecb87409464edc10469ddc9a22a27a99123be49"}, - {file = "Jinja2-2.11.1.tar.gz", hash = "sha256:93187ffbc7808079673ef52771baa950426fd664d3aad1d0fa3e95644360e250"}, + {file = "Jinja2-2.11.2-py2.py3-none-any.whl", hash = "sha256:f0a4641d3cf955324a89c04f3d94663aa4d638abe8f733ecd3582848e1c37035"}, + {file = "Jinja2-2.11.2.tar.gz", hash = "sha256:89aab215427ef59c34ad58735269eb58b1a5808103067f7bb9d5836c651b3bb0"}, ] jsonschema = [ {file = "jsonschema-3.2.0-py2.py3-none-any.whl", hash = "sha256:4e5b3cf8216f577bee9ce139cbe72eca3ea4f292ec60928ff24758ce626cd163"}, diff --git a/pyproject.toml b/pyproject.toml index d3a218d0..9157b66d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,6 @@ python-dateutil = "^2.8" pytz = "^2019.3" requests = "^2.23" yarl = "^1.4.2" -typing-extensions = "^3.7.4" certifi = "^2020.4.5" [tool.poetry.dev-dependencies] @@ -67,7 +66,7 @@ xdoctest = "^0.11.0" towncrier = "^19.2.0" [tool.poetry.extras] -doc = ["sphinx", "sphinx-autodoc-typehints", "sphinx_rtd_theme", "autodocsumm", "tomlkit"] +docs = ["sphinx", "sphinx-autodoc-typehints", "sphinx_rtd_theme", "autodocsumm", "tomlkit"] [tool.black] include = '\.pyi?$' @@ -86,7 +85,7 @@ exclude = ''' [tool.isort] known_first_party = 'robinhood' -known_third_party = ["certifi", "dateutil", "freezegun", "marshmallow", "pytest", "pytz", "requests", "requests_mock", "typing_extensions", "yarl"] +known_third_party = ["certifi", "dateutil", "freezegun", "marshmallow", "pytest", "pytz", "requests", "requests_mock", "yarl"] multi_line_output = 3 lines_after_imports = 2 force_grid_wrap = 0 diff --git a/pyrh/models/__init__.py b/pyrh/models/__init__.py index 7229e58c..7a612ed9 100644 --- a/pyrh/models/__init__.py +++ b/pyrh/models/__init__.py @@ -1,7 +1,17 @@ """pyrh models and schemas.""" -from .oauth import Challenge, OAuth -from .sessionmanager import SessionManager +from .oauth import Challenge, ChallengeSchema, OAuth, OAuthSchema +from .portfolio import PortfolioSchema +from .sessionmanager import SessionManager, SessionManagerSchema -__all__ = ["OAuth", "Challenge", "SessionManager"] +__all__ = [ + "OAuth", + "OAuthSchema", + "Challenge", + "ChallengeSchema", + "SessionManager", + "SessionManagerSchema", + "Portfolio", + "PortfolioSchema", +] diff --git a/pyrh/models/portfolio.py b/pyrh/models/portfolio.py new file mode 100644 index 00000000..5d148ff9 --- /dev/null +++ b/pyrh/models/portfolio.py @@ -0,0 +1,38 @@ +"""Current portfolio.""" + +from marshmallow import fields + +from .base import BaseModel, BaseSchema + + +class Portfolio(BaseModel): + """Robinhood portfolio data class.""" + + pass + + +class PortfolioSchema(BaseSchema): + """Robinhood portfolio schema data loader.""" + + __model__ = Portfolio + + url = fields.URL() + account = fields.URL() + start_date = fields.NaiveDateTime() + market_value = fields.Float() + equity = fields.Float() + extended_hours_market_value = fields.Float() + extended_hours_equity = fields.Float() + extended_hours_portfolio_equity = fields.Float() + last_core_market_value = fields.Float() + last_core_equity = fields.Float() + last_core_portfolio_equity = fields.Float() + excess_margin = fields.Float() + excess_maintenance = fields.Float() + excess_margin_with_uncleared_deposits = fields.Float() + portfolio_equity_previous_close = fields.Float() + adjusted_equity_previous_close = fields.Float() + adjusted_portfolio_equity_previous_close = fields.Float() + withdrawable_amount = fields.Float() + unwithdrawable_deposits = fields.Float() + unwithdrawable_grants = fields.Float() diff --git a/pyrh/models/sessionmanager.py b/pyrh/models/sessionmanager.py index 37f39546..df5b069f 100644 --- a/pyrh/models/sessionmanager.py +++ b/pyrh/models/sessionmanager.py @@ -2,26 +2,15 @@ import uuid from datetime import datetime, timedelta -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Optional, - Union, - cast, - overload, -) +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Union, cast from urllib.request import getproxies import certifi import pytz import requests -from marshmallow import fields, post_load -from requests import Response +from marshmallow import Schema, fields, post_load from requests.exceptions import HTTPError from requests.structures import CaseInsensitiveDict -from typing_extensions import Literal from yarl import URL from pyrh.common import API_BASE, JSON @@ -183,37 +172,7 @@ def login(self, force_refresh: bool = False) -> None: elif self.oauth.is_valid and (self.token_expired or force_refresh): self._refresh_oauth2() - # The following type hints helps mypy determine what the output type to assign based - # on the `return_response` parameter. The same "stub" method approach is used for - # the post method as well. - # https://github.com/python/mypy/issues/8634#issuecomment-609411104 - @overload - def get( - self, - url: Union[str, URL], - params: Optional[Dict[str, Any]] = None, - *, - headers: Optional[CaseInsensitiveDictType] = None, - raise_errors: bool = True, - return_response: Literal[True], - auto_login: bool = True, - ) -> Response: # noqa: D102 - ... - - @overload # noqa: F811 def get( - self, - url: Union[str, URL], - params: Optional[Dict[str, Any]] = None, - *, - headers: Optional[CaseInsensitiveDictType] = None, - raise_errors: bool = True, - return_response: Literal[False] = ..., - auto_login: bool = True, - ) -> JSON: # noqa: D102 - ... - - def get( # noqa: F811 self, url: Union[str, URL], params: Optional[Dict[str, Any]] = None, @@ -222,7 +181,8 @@ def get( # noqa: F811 raise_errors: bool = True, return_response: bool = False, auto_login: bool = True, - ) -> Union[Response, JSON]: + schema: Optional[Schema] = None, + ) -> Any: """Run a wrapped session HTTP GET request. Note: @@ -238,9 +198,12 @@ def get( # noqa: F811 JSON response from the request. auto_login: Whether or not to automatically login on restricted endpoint errors. + schema: An instance of a `marshmallow.Schema` that represents the object + to build. Returns: - The POST response + A JSON dictionary or a constructed object if a schema is passed. If \ + `return_response` is set then a tuple of (response, data) is passed. """ params = {} if params is None else params @@ -255,35 +218,11 @@ def get( # noqa: F811 if raise_errors: res.raise_for_status() - return res if return_response else res.json() + data = res.json() if schema is None else schema.load(res.json()) - @overload - def post( - self, - url: Union[str, URL], - data: Optional[JSON] = None, - *, - headers: Optional[CaseInsensitiveDictType] = None, - raise_errors: bool = True, - return_response: Literal[True], - auto_login: bool = True, - ) -> Response: # noqa: D102 - ... + return (data, res) if return_response else data - @overload # noqa: F811 def post( - self, - url: Union[str, URL], - data: Optional[JSON] = None, - *, - headers: Optional[CaseInsensitiveDictType] = None, - raise_errors: bool = True, - return_response: Literal[False] = ..., - auto_login: bool = True, - ) -> JSON: # noqa: D102 - ... - - def post( # noqa: F811 self, url: Union[str, URL], data: Optional[JSON] = None, @@ -292,7 +231,8 @@ def post( # noqa: F811 raise_errors: bool = True, return_response: bool = False, auto_login: bool = True, - ) -> Union[JSON, Response]: + schema: Optional[Schema] = None, + ) -> Any: """Run a wrapped session HTTP POST request. Note: @@ -308,9 +248,12 @@ def post( # noqa: F811 raise_errors: Whether or not raise errors on POST request. auto_login: Whether or not to automatically login on restricted endpoint errors. + schema: An instance of a `marshmallow.Schema` that represents the object + to build. Returns: - The response or an empty dict if an empty response is returned. + A JSON dictionary or a constructed object if a schema is passed. If \ + `return_response` is set then a tuple of (response, data) is passed. """ res = self.session.post( @@ -327,7 +270,9 @@ def post( # noqa: F811 if raise_errors: res.raise_for_status() - return res if return_response else res.json() + data = res.json() if schema is None else schema.load(res.json()) + + return (data, res) if return_response else data def _configure_manager(self, oauth: OAuth) -> None: """Process an authentication response dictionary. @@ -376,28 +321,30 @@ def _challenge_oauth2(self, oauth: OAuth, oauth_payload: JSON) -> OAuth: challenge_header = CaseInsensitiveDict( {"X-ROBINHOOD-CHALLENGE-RESPONSE-ID": str(oauth.challenge.id)} ) - res = self.post( + oauth_inner, res = self.post( challenge_url, data=challenge_payload, raise_errors=False, headers=challenge_header, auto_login=False, return_response=True, + schema=OAuthSchema(), ) - oauth_inner = OAuthSchema().load(res.json()) if res.status_code == requests.codes.ok: try: - res2 = self.post( - OAUTH, - data=oauth_payload, - headers=challenge_header, - auto_login=False, + # the cast is required for mypy + return cast( + OAuth, + self.post( + OAUTH, + data=oauth_payload, + headers=challenge_header, + auto_login=False, + schema=OAuthSchema(), + ), ) except HTTPError: raise AuthenticationError("Error in finalizing auth token") - else: - oauth = OAuthSchema().load(res2) - return oauth elif oauth_inner.is_challenge and oauth_inner.challenge.can_retry: print("Invalid code entered") return self._challenge_oauth2(oauth, oauth_payload) @@ -424,12 +371,13 @@ def _mfa_oauth2(self, oauth_payload: JSON, attempts: int = 3) -> OAuth: print(f"Input mfa code:") mfa_code = input() oauth_payload["mfa_code"] = mfa_code - res = self.post( + oauth, res = self.post( OAUTH, data=oauth_payload, raise_errors=False, auto_login=False, return_response=True, + schema=OAuthSchema(), ) attempts -= 1 if (res.status_code != requests.codes.ok) and (attempts > 0): @@ -437,7 +385,7 @@ def _mfa_oauth2(self, oauth_payload: JSON, attempts: int = 3) -> OAuth: return self._mfa_oauth2(oauth_payload, attempts) elif res.status_code == requests.codes.ok: # TODO: Write mypy issue on why this needs to be casted? - return cast(OAuth, OAuthSchema().load(res.json())) + return cast(OAuth, oauth) else: raise AuthenticationError("Too many incorrect mfa attempts") @@ -462,16 +410,14 @@ def _login_oauth2(self) -> None: "challenge_type": self.challenge_type, } - res = self.post( + oauth = self.post( OAUTH, data=oauth_payload, raise_errors=False, auto_login=False, - return_response=True, + schema=OAuthSchema(), ) - oauth = OAuthSchema().load(res.json()) - if oauth.is_challenge: oauth = self._challenge_oauth2(oauth, oauth_payload) elif oauth.is_mfa: @@ -507,11 +453,12 @@ def _refresh_oauth2(self) -> None: } self.session.headers.pop("Authorization", None) try: - res = self.post(OAUTH, data=relogin_payload, auto_login=False) + oauth = self.post( + OAUTH, data=relogin_payload, auto_login=False, schema=OAuthSchema() + ) except HTTPError: raise AuthenticationError("Failed to refresh token") - oauth = OAuthSchema().load(res) self._configure_manager(oauth) def logout(self) -> None: