Support unit abbreviations for --max-content-length

ChangeLog.md

@@ -4,6 +4,7 @@ Starting with v1.31.6, this file will contain a record of major features and upd
 
 ## Upcoming
 - Deprecated Python 3.7 support ([Link to PR](https://github.com/aws/graph-notebook/pull/551))
+- Added unit abbreviation support tk `--max-content-length` ([Link to PR](https://github.com/aws/graph-notebook/pull/553))
 
 ## Release 4.0.2 (Dec 14, 2023)
 - Fixed `neptune_ml_utils` imports in `03-Neptune-ML` samples ([Link to PR](https://github.com/aws/graph-notebook/pull/546))

src/graph_notebook/magics/graph_magic.py

@@ -14,6 +14,7 @@
 import os
 import uuid
 import ast
+import re
 from ipyfilechooser import FileChooser
 from enum import Enum
 from copy import copy
@@ -155,6 +156,9 @@
                                           MEDIA_TYPE_NTRIPLES_TEXT, MEDIA_TYPE_TURTLE, MEDIA_TYPE_N3, MEDIA_TYPE_TRIX,
                                           MEDIA_TYPE_TRIG, MEDIA_TYPE_RDF4J_BINARY]
 
+byte_units = {'B': 1, 'KB': 1024, 'MB': 1024**2, 'GB': 1024**3, 'TB': 1024**4}
+
+
 class QueryMode(Enum):
     DEFAULT = 'query'
     EXPLAIN = 'explain'
@@ -269,6 +273,17 @@ def process_statistics_400(is_summary: bool, response):
     print(f"\nFull response: {bad_request_res}")
 
 
+def mcl_to_bytes(mcl):
+    using_abb = re.match(r'(\d+)([A-Za-z]+)?', mcl, re.IGNORECASE)
+    if using_abb:
+        num, unit = using_abb.groups()
+        unit = unit.upper() if unit else 'B'
+        if unit in byte_units:
+            mcl_bytes = int(num) * byte_units[unit]
+            return mcl_bytes
+    return byte_units['MB'] * 10
+
+
 # TODO: refactor large magic commands into their own modules like what we do with %neptune_ml
 # noinspection PyTypeChecker
 @magics_class
@@ -866,9 +881,10 @@ def gremlin(self, line, cell, local_ns: dict = None):
                             help="Display the entire output without a scroll bar.")
         parser.add_argument('--hide-index', action='store_true', default=False,
                             help="Hide the index column numbers when displaying the results.")
-        parser.add_argument('-mcl', '--max-content-length', type=int, default=10*1024*1024,
+        parser.add_argument('-mcl', '--max-content-length', type=str, default='',
                             help="Specifies maximum size (in bytes) of results that can be returned to the "
-                                 "GremlinPython client. Default is 10MB")
+                                 "GremlinPython client. Abbreviated memory units (ex.'50MB') are accepted. "
+                                 "Default is 10MB")
 
         args = parser.parse_args(line.split())
         mode = str_to_query_mode(args.query_mode)
@@ -894,7 +910,8 @@ def gremlin(self, line, cell, local_ns: dict = None):
             first_tab_output = widgets.Output(layout=gremlin_layout)
             children.append(first_tab_output)
 
-        transport_args = {'max_content_length': args.max_content_length}
+        mcl_bytes = mcl_to_bytes(args.max_content_length)
+        transport_args = {'max_content_length': mcl_bytes}
 
         if mode == QueryMode.EXPLAIN:
             res = self.client.gremlin_explain(cell,