Skip to content

Commit

Permalink
Addressed review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
Sachin Saharan committed Jan 13, 2025
1 parent b377c95 commit ed1953c
Show file tree
Hide file tree
Showing 2 changed files with 114 additions and 23 deletions.
51 changes: 37 additions & 14 deletions lib/python/pyflyby/_saveframe_reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,13 @@ class SaveframeReader:
>> reader.get_variables('var2')
var2_value3
>> reader.get_variables('var1', 'var3')
{2: {'var1': var1_value2, 'var3': var3_value2},
4: {'var1': var1_value4}, 5: {'var3': var3_value5}}
>> reader.get_variables('var1', 'var3', frame_idx=2)
{'var1': var1_value2, 'var3': var3_value2}
>> reader.get_variables(['var1', 'var3'])
{2: {'var1': var1_value2, 'var3': var3_value2},
4: {'var1': var1_value4}, 5: {'var3': var3_value5}}
Expand Down Expand Up @@ -185,7 +192,7 @@ def variables(self):
return frame_idx_to_variables_map


def get_metadata(self, metadata, frame_idx=None):
def get_metadata(self, metadata, *, frame_idx=None):
"""
Retrieve the value of a specific metadata field.
Expand Down Expand Up @@ -280,7 +287,7 @@ def get_metadata(self, metadata, frame_idx=None):
f"are: {allowed_frame_idx}.")


def get_variables(self, variables, frame_idx=None):
def get_variables(self, *variables, frame_idx=None):
"""
Retrieve the value of local variable(s) from specific frames.
Expand All @@ -299,15 +306,24 @@ def get_variables(self, variables, frame_idx=None):
>> reader.get_variables('var2')
var2_value3 # 'var2' is only present in frame 3
>> reader.get_variables(['var1', 'var3'])
>> reader.get_variables('var1', 'var3')
{2: {'var1': var1_value2, 'var3': var3_value2},
4: {'var1': var1_value4}, 5: {'var3': var3_value5}}
>> reader.get_variables('var1', 'var3', frame_idx=2)
{'var1': var1_value2, 'var3': var3_value2}
>> reader.get_variables(['var1', 'var3'], frame_idx=2)
{'var1': var1_value2, 'var3': var3_value2}
:param variables:
A string or a list of variable names for which to retrieve the values.
One or more variable names for which to retrieve the values. You can
pass a single variable name as a string, a list of variable names, or
multiple variable names as separate argument. Example:
- get_variables('var1')
- get_variables(['var1', 'var2', 'var3', ...])
- get_variables('var1', 'var2', 'var3', ...)
:param frame_idx:
The index of the frame from which to retrieve the value(s) of the
variable(s). Default is None, which means values from all frames are
Expand Down Expand Up @@ -336,14 +352,22 @@ def get_variables(self, variables, frame_idx=None):
an error is raised.
"""
# Sanity checks.
if not isinstance(variables, (str, list, tuple)):
raise TypeError(
f"'variables' must either be a string or a list/tuple. "
f"Got '{type(variables).__name__}'.")
if isinstance(variables, (list, tuple)) and len(variables) == 0:
if len(variables) == 1 and isinstance(variables[0], (list, tuple)):
variables = tuple(variables[0])
if len(variables) == 0:
raise ValueError("No 'variables' passed.")
if isinstance(variables, str):
variables = [variables]
for variable in variables:
if not isinstance(variable, str):
raise TypeError(
f"Invalid variable name: {variable}. Each variable name "
f"must be of type string, not {type(variable).__name__}.")

def _get_variable_value_on_unpickle_error(err):
"""
Get variable's value when it fails to unpickle due to error ``err``.
"""
return f"Can't un-pickle the variable. Error: {err}"

# frame_idx is not passed.
if frame_idx is None:
frame_idx_to_variables_map = {}
Expand All @@ -362,8 +386,7 @@ def get_variables(self, variables, frame_idx=None):
logging.warning(
"Can't un-pickle the value of variable %a for frame "
"%a. Error: %s", variable, key_item, err)
variable_value = (
f"Can't un-pickle the variable. Error: {err}")
variable_value = _get_variable_value_on_unpickle_error(err)
if len(variables) == 1:
# Single variable is queried.
frame_idx_to_variables_map[key_item] = variable_value
Expand Down Expand Up @@ -409,7 +432,7 @@ def get_variables(self, variables, frame_idx=None):
"Can't un-pickle the value of variable %a for frame "
"%a. Error: %s", variable, frame_idx, err)
if len(variables) > 1:
variable_value = f"Can't un-pickle the variable. Error: {err}"
variable_value = _get_variable_value_on_unpickle_error(err)
if len(variables) == 1:
# Single variable is queried. Directly return the value.
return variable_value
Expand Down
86 changes: 77 additions & 9 deletions tests/test_saveframe_reader.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -336,7 +336,7 @@ def test_get_metadata_2(tmpdir):
expected = {3: 9, 4: 14}
assert result == expected

result = reader.get_metadata("lineno", 4)
result = reader.get_metadata("lineno", frame_idx=4)
expected = 14
assert result == expected

Expand All @@ -350,7 +350,7 @@ def test_get_metadata_3(tmpdir):
expected = {1: 'func3', 2: 'func2', 3: 'func2', 4: 'func1', 5: 'init_func1'}
assert result == expected

result = reader.get_metadata("function_name", 5)
result = reader.get_metadata("function_name", frame_idx=5)
expected = 'init_func1'
assert result == expected

Expand Down Expand Up @@ -393,7 +393,7 @@ def test_get_metadata_5(tmpdir):
assert name == frame_idx_to_info[key]['name']
assert qualname == frame_idx_to_info[key]['qualname']

result = reader.get_metadata("function_object", 2)
result = reader.get_metadata("function_object", frame_idx=2)
if not isinstance(result, str):
assert result.__qualname__ == get_func2_qualname()

Expand All @@ -408,7 +408,7 @@ def test_get_metadata_6(tmpdir):
3: f'{pkg_name}.mod1', 4: f'{pkg_name}.mod1', 5: pkg_name}
assert result == expected

result = reader.get_metadata("module_name", 4)
result = reader.get_metadata("module_name", frame_idx=4)
assert result == f'{pkg_name}.mod1'


Expand All @@ -422,10 +422,10 @@ def test_get_metadata_7(tmpdir):
3: 'obj.func2()', 4: 'func2()', 5: 'func1()'}
assert result == expected

result = reader.get_metadata("code", 1)
result = reader.get_metadata("code", frame_idx=1)
assert result == 'raise ValueError("Error is raised")'

result = reader.get_metadata("code", 3)
result = reader.get_metadata("code", frame_idx=3)
assert result == 'obj.func2()'


Expand Down Expand Up @@ -660,6 +660,49 @@ def test_get_variables_9(tmpdir):
assert result == expected


def test_get_variables_10(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
reader = SaveframeReader(filename)

result = reader.get_variables('var1', 'var2')
expected = {1: {'var1': [4, 'foo', 2.4], 'var2': 'blah'},
2: {'var1': 'foo', 'var2': (4, 9, 10)},
3: {'var1': 'func2', 'var2': 34}, 4: {'var1': [4, 5, 2]},
5: {'var1': 3, 'var2': 'blah'}}
assert result == expected


def test_get_variables_11(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
reader = SaveframeReader(filename)

result = reader.get_variables('func1_var2', 'func3_var3', 'var3')
expected = {1: {'func3_var3': True}, 4: {'func1_var2': 4.56}}
assert result == expected


def test_get_variables_12(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames='mod1.py::')
reader = SaveframeReader(filename)

result = reader.get_variables('func1_var2', 'func3_var3', 'var3')
expected = {'func1_var2': 4.56}
assert result == expected


def test_get_variables_13(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
reader = SaveframeReader(filename)

result = reader.get_variables('var1', 'var2', 'func1_var2', frame_idx=3)
expected = {'var1': 'func2', 'var2': 34}
assert result == expected


def test_get_variables_invalid_variable_1(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
Expand All @@ -668,7 +711,7 @@ def test_get_variables_invalid_variable_1(tmpdir):
with pytest.raises(ValueError) as err:
reader.get_variables('foo')

expected = "Local variable(s) ['foo'] not found in any of the saved frames."
expected = "Local variable(s) ('foo',) not found in any of the saved frames."
assert str(err.value) == expected


Expand All @@ -680,7 +723,7 @@ def test_get_variables_invalid_variable_2(tmpdir):
with pytest.raises(ValueError) as err:
reader.get_variables('var2', frame_idx=4)

expected = "Local variable(s) ['var2'] not found in the frame 4"
expected = "Local variable(s) ('var2',) not found in the frame 4"
assert str(err.value) == expected


Expand All @@ -692,7 +735,32 @@ def test_get_variables_invalid_variable_3(tmpdir):
with pytest.raises(ValueError) as err:
reader.get_variables(['var2', 'var5'], frame_idx=4)

expected = "Local variable(s) ['var2', 'var5'] not found in the frame 4"
expected = "Local variable(s) ('var2', 'var5') not found in the frame 4"
assert str(err.value) == expected


def test_get_variables_invalid_variable_4(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
reader = SaveframeReader(filename)

with pytest.raises(ValueError) as err:
reader.get_variables()

expected = "No 'variables' passed."
assert str(err.value) == expected


def test_get_variables_invalid_variable_5(tmpdir):
pkg_name = create_pkg(tmpdir)
filename = call_saveframe(pkg_name, tmpdir, frames=5)
reader = SaveframeReader(filename)

with pytest.raises(TypeError) as err:
reader.get_variables('var1', 2)

expected = ("Invalid variable name: 2. Each variable name must be of type "
"string, not int.")
assert str(err.value) == expected


Expand Down

0 comments on commit ed1953c

Please sign in to comment.