Merging upstream version 1.12.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>

7 files changed, 119 insertions, 33 deletions

diff --git a/iredis/__init__.py b/iredis/__init__.py
index b518f6e..438a38d 100644
--- a/iredis/__init__.py
+++ b/iredis/__init__.py
@@ -1 +1 @@
-__version__ = "1.12.0"
+__version__ = "1.12.1"
diff --git a/iredis/client.py b/iredis/client.py
index 1b1bf38..6555ff1 100644
--- a/iredis/client.py
+++ b/iredis/client.py
@@ -55,25 +55,35 @@ class Client:
 
     def __init__(
         self,
-        host=None,
-        port=None,
+        host="127.0.0.1",
+        port=6379,
         db=0,
         password=None,
         path=None,
         scheme="redis",
         username=None,
         client_name=None,
+        prompt=None,
     ):
         self.host = host
         self.port = port
         self.db = db
         self.path = path
-        # FIXME username is not using...
         self.username = username
         self.client_name = client_name
         self.scheme = scheme
         self.password = password
 
+        # cli args --prompt will overwrite the prompt in iredisrc config file
+        self.prompt = ""
+        if config.prompt:
+            self.prompt = config.prompt
+        if prompt:
+            self.prompt = prompt
+
+        self.client_id = None
+        self.client_addr = None
+
         self.build_connection()
 
         # all command upper case
@@ -93,6 +103,13 @@ class Client:
         else:
             config.no_version_reason = "--no-info flag activated"
 
+        if self.prompt and "client_addr" in self.prompt:
+            self.client_addr = ":".join(
+                str(x) for x in self.connection._sock.getsockname()
+            )
+        if self.prompt and "client_id" in self.prompt:
+            self.client_id = str(self.execute("CLIENT ID"))
+
         if config.version and re.match(r"([\d\.]+)", config.version):
             self.auth_compat(config.version)
 
@@ -131,6 +148,11 @@ class Client:
                 "socket_keepalive": config.socket_keepalive,
                 "client_name": client_name,
             }
+
+            # if username is set without setting paswword, password will be ignored
+            if password:
+                connection_kwargs["username"] = username
+
             if scheme == "rediss":
                 connection_class = SSLConnection
             else:
@@ -141,6 +163,7 @@ class Client:
                 "password": password,
                 "path": path,
                 "client_name": client_name,
+                "username": username,
             }
             connection_class = UnixDomainSocketConnection
 
@@ -150,7 +173,8 @@ class Client:
             connection_kwargs["encoding_errors"] = "replace"
 
         logger.debug(
-            f"connection_class={connection_class}, connection_kwargs={connection_kwargs}"
+            f"connection_class={connection_class},"
+            f" connection_kwargs={connection_kwargs}"
         )
 
         return connection_class(**connection_kwargs)
@@ -191,6 +215,18 @@ class Client:
         config.version = version
 
     def __str__(self):
+        if self.prompt:  # not None and not empty
+            return self.prompt.format(
+                client_name=self.client_name,
+                db=self.db,
+                host=self.host,
+                path=self.path,
+                port=self.port,
+                username=self.username,
+                client_addr=self.client_addr,
+                client_id=self.client_id,
+            )
+
         if self.scheme == "unix":
             prompt = f"redis {self.path}"
         else:
@@ -198,7 +234,8 @@ class Client:
 
         if self.db:
             prompt = f"{prompt}[{self.db}]"
-        return prompt
+
+        return f"{prompt}> "
 
     def client_execute_command(self, command_name, *args):
         command = command_name.upper()
@@ -222,7 +259,8 @@ class Client:
         Here we retry once for ConnectionError.
         """
         logger.info(
-            f"execute by connection: connection={connection}, name={command_name}, {args}, {options}"
+            f"execute by connection: connection={connection}, name={command_name},"
+            f" {args}, {options}"
         )
         retry_times = config.retry_times  # FIXME configurable
         last_error = None
@@ -248,7 +286,7 @@ class Client:
                 last_error = e
                 retry_times -= 1
                 need_refresh_connection = True
-            except (ResponseError) as e:
+            except ResponseError as e:
                 response_message = str(e)
                 if response_message.startswith("MOVED"):
                     return self.reissue_with_redirect(
diff --git a/iredis/config.py b/iredis/config.py
index fefccdb..a62101c 100644
--- a/iredis/config.py
+++ b/iredis/config.py
@@ -63,6 +63,8 @@ class Config:
         self.withscores = False
         self.version = "Unknown"
 
+        self.prompt = None
+
     def __setter__(self, name, value):
         # for every time start a transaction
         # clear the queued commands first
@@ -126,5 +128,6 @@ def load_config_files(iredisrc):
     config.shell = config_obj["main"].as_bool("shell")
     config.pager = config_obj["main"].get("pager")
     config.enable_pager = config_obj["main"].as_bool("enable_pager")
+    config.prompt = config_obj["main"].get("prompt")
 
     return config_obj
diff --git a/iredis/data/command_syntax.csv b/iredis/data/command_syntax.csv
index 477c15c..fbaf47b 100644
--- a/iredis/data/command_syntax.csv
+++ b/iredis/data/command_syntax.csv
@@ -87,6 +87,7 @@ hash,HKEYS,command_key,command_hkeys
 hash,HLEN,command_key,render_int
 hash,HMGET,command_key_fields,render_list
 hash,HMSET,command_key_fieldvalues,render_bulk_string
+hash,HRANDFIELD,command_key_count_withvalues,render_list_or_string
 hash,HSCAN,command_key_cursor_match_pattern_count,command_hscan
 hash,HSET,command_key_field_value,render_int
 hash,HSETNX,command_key_field_value,render_int
@@ -103,14 +104,14 @@ list,LINDEX,command_key_position,render_bulk_string
 list,LINSERT,command_key_positionchoice_pivot_value,render_int
 list,LLEN,command_key,render_int
 list,LPOS,command_lpos,render_list_or_string
-list,LPOP,command_key,render_bulk_string
+list,LPOP,command_key,render_list_or_string
 list,LPUSH,command_key_values,render_int
 list,LPUSHX,command_key_values,render_int
 list,LRANGE,command_key_start_end,render_list
 list,LREM,command_key_position_value,render_int
 list,LSET,command_key_position_value,render_simple_string
 list,LTRIM,command_key_start_end,render_simple_string
-list,RPOP,command_key,render_bulk_string
+list,RPOP,command_key,render_list_or_string
 list,RPOPLPUSH,command_key_newkey,render_bulk_string
 list,RPUSH,command_key_values,render_int
 list,RPUSHX,command_key_value,render_int
diff --git a/iredis/data/iredisrc b/iredis/data/iredisrc
index bf79e90..6db02a7 100644
--- a/iredis/data/iredisrc
+++ b/iredis/data/iredisrc
@@ -58,6 +58,21 @@ warning = True
 # eg. ~/.iredis.log
 log_location =
 
+# You can change the prompt str, if left blank, the default prompt would be:
+# 127.0.0.1:6379> 
+# which is rendered by "{host}:{port}[{db}]> "
+# supported interpolations:
+# {client_name}
+# {db}
+# {host}
+# {path}
+# {port}
+# {username}
+# {client_addr}
+# {client_id}
+# The prompt string uses python string format engine
+prompt = 
+
 # History file location
 history_location = ~/.iredis_history
 
diff --git a/iredis/entry.py b/iredis/entry.py
index 2a9dc05..9a07956 100644
--- a/iredis/entry.py
+++ b/iredis/entry.py
@@ -123,22 +123,22 @@ def write_result(text, max_height=None):
 
 class Rainbow:
     color = [
-        ("#cc2244"),
-        ("#bb4444"),
-        ("#996644"),
-        ("#cc8844"),
-        ("#ccaa44"),
-        ("#bbaa44"),
-        ("#99aa44"),
-        ("#778844"),
-        ("#55aa44"),
-        ("#33aa44"),
-        ("#11aa44"),
-        ("#11aa66"),
-        ("#11aa88"),
-        ("#11aaaa"),
-        ("#11aacc"),
-        ("#11aaee"),
+        "#cc2244",
+        "#bb4444",
+        "#996644",
+        "#cc8844",
+        "#ccaa44",
+        "#bbaa44",
+        "#99aa44",
+        "#778844",
+        "#55aa44",
+        "#33aa44",
+        "#11aa44",
+        "#11aa66",
+        "#11aa88",
+        "#11aaaa",
+        "#11aacc",
+        "#11aaee",
     ]
 
     def __init__(self):
@@ -160,8 +160,7 @@ class Rainbow:
 
 
 def prompt_message(client):
-    # TODO custom prompt
-    text = "{hostname}> ".format(hostname=str(client))
+    text = str(client)
     if config.rainbow:
         return list(zip(Rainbow(), text))
     return text
@@ -248,8 +247,11 @@ PAGER_HELP = """Using pager when output is too tall for your window, default to
 @click.option(
     "-s", "--socket", default=None, help="Server socket (overrides hostname and port)."
 )
+@click.option("-n", help="Database number.(overwrites dsn/url's db number)", default=0)
 @click.option(
-    "-n", help="Database number.(overwrites dsn/url's db number)", default=None
+    "-u",
+    "--username",
+    help="User name used to auth, will be ignore for redis version < 6.",
 )
 @click.option("-a", "--password", help="Password to use when connecting to the server.")
 @click.option("--url", default=None, envvar="IREDIS_URL", help=URL_HELP)
@@ -271,6 +273,14 @@ PAGER_HELP = """Using pager when output is too tall for your window, default to
 @click.option("--rainbow/--no-rainbow", default=None, is_flag=True, help=RAINBOW)
 @click.option("--shell/--no-shell", default=None, is_flag=True, help=SHELL)
 @click.option("--pager/--no-pager", default=None, is_flag=True, help=PAGER_HELP)
+@click.option(
+    "--prompt",
+    default=None,
+    help=(
+        "Prompt format (supported interpolations: {client_name}, {db}, {host}, {path},"
+        " {port}, {username}, {client_addr}, {client_id})."
+    ),
+)
 @click.version_option()
 @click.argument("cmd", nargs=-1)
 def gather_args(
@@ -278,6 +288,7 @@ def gather_args(
     h,
     p,
     n,
+    username,
     password,
     client_name,
     newbie,
@@ -291,6 +302,7 @@ def gather_args(
     socket,
     shell,
     pager,
+    prompt,
 ):
     """
     IRedis: Interactive Redis
@@ -311,9 +323,9 @@ def gather_args(
     load_config_files(iredisrc)
     setup_log()
     logger.info(
-        f"[commandline args] host={h}, port={p}, db={n}, newbie={newbie}, "
-        f"iredisrc={iredisrc}, decode={decode}, raw={raw}, "
-        f"cmd={cmd}, rainbow={rainbow}."
+        f"[commandline args] host={h}, port={p}, db={n}, user={username},"
+        f" newbie={newbie}, iredisrc={iredisrc}, decode={decode}, raw={raw}, cmd={cmd},"
+        f" rainbow={rainbow}."
     )
     # raw config
     if raw is not None:
@@ -368,8 +380,10 @@ def create_client(params):
     host = params["h"]
     port = params["p"]
     db = params["n"]
+    username = params["username"]
     password = params["password"]
     client_name = params["client_name"]
+    prompt = params["prompt"]
 
     dsn_from_url = None
     dsn = params["dsn"]
@@ -390,17 +404,26 @@ def create_client(params):
             scheme=dsn_from_url.scheme,
             username=dsn_from_url.username,
             client_name=client_name,
+            prompt=prompt,
         )
     if params["socket"]:
         return Client(
             scheme="unix",
             path=params["socket"],
             db=db,
+            username=username,
             password=password,
             client_name=client_name,
+            prompt=prompt,
         )
     return Client(
-        host=host, port=port, db=db, password=password, client_name=client_name
+        host=host,
+        port=port,
+        db=db,
+        username=username,
+        password=password,
+        client_name=client_name,
+        prompt=prompt,
     )
 
 
diff --git a/iredis/redis_grammar.py b/iredis/redis_grammar.py
index 4e38540..b9e2ea8 100644
--- a/iredis/redis_grammar.py
+++ b/iredis/redis_grammar.py
@@ -16,6 +16,7 @@ logger = logging.getLogger(__name__)
 CONST = {
     "failoverchoice": "TAKEOVER FORCE",
     "withscores": "WITHSCORES",
+    "withvalues_const": "WITHVALUES",
     "limit": "LIMIT",
     "expiration": "EX PX",
     "exat_const": "EXAT",
@@ -362,6 +363,7 @@ TIMEOUT_CONST = rf"(?P<timeout_const>{c('timeout_const')})"
 ABORT_CONST = rf"(?P<abort_const>{c('abort_const')})"
 PXAT_CONST = rf"(?P<pxat_const>{c('pxat_const')})"
 EXAT_CONST = rf"(?P<exat_const>{c('exat_const')})"
+WITHVALUES_CONST = rf"(?P<withvalues_const>{c('withvalues_const')})"
 
 command_grammar = compile(COMMAND)
 
@@ -660,6 +662,10 @@ GRAMMAR = {
             (\s+ {EXAT_CONST} \s+ {TIMESTAMP})
         )?
         \s*""",
+    "command_key_count_withvalues": rf"""
+        \s+ {KEY}
+        (\s+ {COUNT} (\s+ {WITHVALUES_CONST})?)?
+        \s*""",
 }
 
 pipeline = r"(?P<shellcommand>\|.*)?"