From 26bb0b67fc02f9e98dd25924cb4a885313fa43ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Timoth=C3=A9e=20Mazzucotelli?= <pawamoy@pm.me>
Date: Tue, 7 Apr 2020 17:31:36 +0200
Subject: [PATCH] feat: Add a timeout to client's requests

Reference: issue #52.
---
 src/aria2p/cli.py       | 19 +++++++++++++++++--
 src/aria2p/client.py    |  9 +++++++--
 src/aria2p/interface.py |  8 ++++++--
 3 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/src/aria2p/cli.py b/src/aria2p/cli.py
index 505602b..baeb5bc 100644
--- a/src/aria2p/cli.py
+++ b/src/aria2p/cli.py
@@ -29,7 +29,7 @@
 from aria2p import Download, enable_logger
 
 from .api import API
-from .client import DEFAULT_HOST, DEFAULT_PORT, Client, ClientException
+from .client import DEFAULT_HOST, DEFAULT_PORT, DEFAULT_TIMEOUT, Client, ClientException
 from .utils import get_version
 
 try:
@@ -61,7 +61,14 @@ def main(args: Optional[List[str]] = None) -> int:
     check_args(parser, args)
 
     logger.debug("Instantiating API")
-    api = API(Client(host=kwargs.pop("host"), port=kwargs.pop("port"), secret=kwargs.pop("secret")))
+    api = API(
+        Client(
+            host=kwargs.pop("host"),
+            port=kwargs.pop("port"),
+            secret=kwargs.pop("secret"),
+            timeout=kwargs.pop("client_timeout"),
+        )
+    )
 
     logger.info(f"API instantiated: {api!r}")
 
@@ -166,6 +173,14 @@ def get_parser() -> argparse.ArgumentParser:
     global_options.add_argument(
         "-P", "--log-path", dest="log_path", default=None, help="Log path to use. Can be a directory or a file."
     )
+    global_options.add_argument(
+        "-T",
+        "--client-timeout",
+        dest="client_timeout",
+        default=DEFAULT_TIMEOUT,
+        type=float,
+        help=f"Timeout in seconds for requests to the remote server. Floats supported. Default: {DEFAULT_TIMEOUT}.",
+    )
 
     # ========= SUBPARSERS ========= #
     subparsers = parser.add_subparsers(dest="subcommand", title="Commands", metavar="", prog="aria2p")
diff --git a/src/aria2p/client.py b/src/aria2p/client.py
index 3906484..2e067b3 100644
--- a/src/aria2p/client.py
+++ b/src/aria2p/client.py
@@ -15,6 +15,7 @@
 DEFAULT_ID = -1
 DEFAULT_HOST = "http://localhost"
 DEFAULT_PORT = 6800
+DEFAULT_TIMEOUT: float = 60.0
 
 JSONRPC_PARSER_ERROR = -32700
 JSONRPC_INVALID_REQUEST = -32600
@@ -169,7 +170,9 @@ class Client:
         LIST_NOTIFICATIONS,
     ]
 
-    def __init__(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, secret: str = "") -> None:  # nosec
+    def __init__(  # nosec
+        self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, secret: str = "", timeout: float = DEFAULT_TIMEOUT,
+    ) -> None:
         """
         Initialization method.
 
@@ -177,12 +180,14 @@ def __init__(self, host: str = DEFAULT_HOST, port: int = DEFAULT_PORT, secret: s
             host: the remote process address.
             port: the remote process port.
             secret: the secret token.
+            timeout: the timeout to use for requests towards the remote server.
         """
         host = host.rstrip("/")
 
         self.host = host
         self.port = port
         self.secret = secret
+        self.timeout = timeout
         self.listening = False
 
     def __str__(self):
@@ -321,7 +326,7 @@ def post(self, payload: dict) -> dict:
         Returns:
             The answer from the server, as a Python dictionary.
         """
-        return requests.post(self.server, data=payload).json()
+        return requests.post(self.server, data=payload, timeout=self.timeout).json()
 
     @staticmethod
     def response_as_exception(response: dict) -> ClientException:
diff --git a/src/aria2p/interface.py b/src/aria2p/interface.py
index 328bc9e..5eee66a 100644
--- a/src/aria2p/interface.py
+++ b/src/aria2p/interface.py
@@ -22,6 +22,7 @@
 from collections import defaultdict
 from pathlib import Path
 
+import requests
 from asciimatics.event import KeyboardEvent, MouseEvent
 from asciimatics.screen import ManagedScreen, Screen
 from loguru import logger
@@ -860,8 +861,11 @@ def get_data(self):
 
     def update_data(self):
         """Set the interface data and rows contents."""
-        self.data = self.get_data()
-        self.sort_data()
+        try:
+            self.data = self.get_data()
+            self.sort_data()
+        except requests.exceptions.Timeout:
+            logger.debug("Request timeout")
 
     def sort_data(self):
         """Sort data according to interface state."""