diff --git a/ChangeLog.md b/ChangeLog.md index 17be8a1e..735127d1 100644 --- a/ChangeLog.md +++ b/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)) diff --git a/src/graph_notebook/magics/graph_magic.py b/src/graph_notebook/magics/graph_magic.py index 8bf1e989..c1e0d717 100644 --- a/src/graph_notebook/magics/graph_magic.py +++ b/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,