Skip to content

Commit

Permalink
add data-tree list-attributes (#1025)
Browse files Browse the repository at this point in the history
* add data-tree list-attributes

* fix style

* fix style

* add test skip for older versions

* add to_dict() & docstring examples

* add newline

* fix style

* fix gdc test

* fix style check

* rename test
  • Loading branch information
jose-henriquezroa authored Jul 19, 2023
1 parent 01e38f9 commit 99e8b4f
Show file tree
Hide file tree
Showing 2 changed files with 130 additions and 0 deletions.
85 changes: 85 additions & 0 deletions src/ansys/dpf/core/data_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,6 +482,91 @@ def get_as(self, name, type_to_return=types.string):
out = DataTree(data_tree=obj, server=self._server)
return out

def get_attribute_names(self):
"""
Returns a list of defined attribute names.
Returns
-------
list[str]
Examples
--------
>>> from ansys.dpf import core as dpf
>>> data_tree = dpf.DataTree()
>>> data_tree.add(id=3, qualities=["nice", "funny"], name="George")
>>> data_tree.get_attribute_names()
['id', 'name', 'qualities']
"""
coll_obj = collection.StringCollection(
collection=self._api.dpf_data_tree_get_available_attributes_names_in_string_collection(
self
),
server=self._server,
)

return coll_obj.get_integral_entries()

def get_sub_tree_names(self):
"""
Returns a list of defined sub-tree names.
Returns
-------
list[str]
Examples
--------
>>> from ansys.dpf import core as dpf
>>> data_tree = dpf.DataTree()
>>> first_subtree = dpf.DataTree()
>>> second_subtree = dpf.DataTree()
>>> data_tree.add(first=first_subtree, second=second_subtree)
>>> data_tree.get_sub_tree_names()
['first', 'second']
"""
coll_obj = collection.StringCollection(
collection=self._api.dpf_data_tree_get_available_sub_tree_names_in_string_collection(
self
),
server=self._server,
)

return coll_obj.get_integral_entries()

def __to_dict(self, dic):
for attribute_name in self.get_attribute_names():
dic[attribute_name] = self.get_as(attribute_name)

for sub_tree_name in self.get_sub_tree_names():
sub_tree = self.get_as(sub_tree_name, types.data_tree)
sub_dic = {}
sub_tree.__to_dict(sub_dic)
dic[sub_tree_name] = sub_dic

def to_dict(self):
"""
Returns a dictionary representation of the DataTree
Returns
-------
dict
Examples
--------
>>> from ansys.dpf import core as dpf
>>> data_tree = dpf.DataTree()
>>> sub = dpf.DataTree()
>>> sub.add(str="hello world")
>>> data_tree.add(id=3, sub_tree=sub)
>>> data_tree.to_dict()
{'id': '3', 'sub_tree': {'str': 'hello world'}}
"""
dic = {}
self.__to_dict(dic)

return dic

def __setattr__(self, key, value):
if key == "_common_keys" or key in self._common_keys:
return super.__setattr__(self, key, value)
Expand Down
45 changes: 45 additions & 0 deletions tests/test_data_tree.py
Original file line number Diff line number Diff line change
Expand Up @@ -329,3 +329,48 @@ def test_unsupported_types_data_tree(server_type):
data_tree.add(data1=[[1]])
with pytest.raises(TypeError):
data_tree.add(data1=(1, 2))


@pytest.mark.skipif(
not conftest.SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0, reason="Available for servers >=7.0"
)
def test_list_attributes_data_tree(server_type):
data_tree = dpf.DataTree(server=server_type)
with data_tree.to_fill() as to_fill:
to_fill.int = 1
to_fill.double = 1.0
to_fill.string = "hello"
to_fill.list_int = [1, 2]
to_fill.list_double = [1.5, 2.5]
to_fill.add(list_string=["hello", "bye"])

attributes = data_tree.get_attribute_names()

assert ["double", "int", "list_double", "list_int", "list_string", "string"] == attributes


@pytest.mark.skipif(
not conftest.SERVERS_VERSION_GREATER_THAN_OR_EQUAL_TO_7_0, reason="Available for servers >=7.0"
)
def test_list_attributes_recursive_data_tree(server_type):
data_tree = dpf.DataTree(server=server_type)
with data_tree.to_fill() as to_fill:
to_fill.attribute01 = 1
sub_tree01 = dpf.DataTree(server=server_type)
with sub_tree01.to_fill() as to_fill01:
to_fill01.attribute02 = 2
to_fill.sub_tree01 = sub_tree01
sub_tree02 = dpf.DataTree(server=server_type)
to_fill.sub_tree02 = sub_tree02

attributes = data_tree.get_attribute_names()
sub_trees = data_tree.get_sub_tree_names()

assert attributes == ["attribute01"]
assert sub_trees == ["sub_tree01", "sub_tree02"]

dic = data_tree.to_dict()

assert ["attribute01", "sub_tree01", "sub_tree02"] == list(dic.keys())
assert {"attribute02": "2"} == dic["sub_tree01"]
assert {} == dic["sub_tree02"]

0 comments on commit 99e8b4f

Please sign in to comment.