[yt-svn] commit/yt: jzuhone: Merged in ngoldbaum/yt (pull request #2459)

[commits-noreply at bitbucket.org]

1 new commit in yt:

https://bitbucket.org/yt_analysis/yt/commits/c5bdb1f946cd/
Changeset:   c5bdb1f946cd
Branch:      yt
User:        jzuhone
Date:        2016-12-06 15:33:24+00:00
Summary:     Merged in ngoldbaum/yt (pull request #2459)

Implement iterator protocol for FieldTypeContainer and FieldNameContainer
Affected #:  3 files

diff -r bcdf7ee3c05bc4c96372fc3bfa5a581595c09bb2 -r c5bdb1f946cd9a9233b36a13a758840522b78914 doc/source/analyzing/fields.rst
--- a/doc/source/analyzing/fields.rst
+++ b/doc/source/analyzing/fields.rst
@@ -64,6 +64,24 @@
 You can use this to easily explore available fields, particularly through
 tab-completion in Jupyter/IPython.
 
+It's also possible to iterate over the list of fields associated with each
+field type. For example, to print all of the ``'gas'`` fields, one might do:
+
+.. code-block:: python
+
+   for field in ds.fields.gas:
+       print(field)
+
+You can also check if a given field is associated with a field type using
+standard python syntax:
+
+.. code-block:: python
+
+   # these examples evaluate to True for a dataset that has ('gas', 'density')
+   'density' in ds.fields.gas
+   ('gas', 'density') in ds.fields.gas
+   ds.fields.gas.density in ds.fields.gas
+
 For a more programmatic method of accessing fields, you can utilize the
 ``ds.field_list``, ``ds.derived_field_list`` and some accessor methods to gain
 information about fields.  The full list of fields available for a dataset can

diff -r bcdf7ee3c05bc4c96372fc3bfa5a581595c09bb2 -r c5bdb1f946cd9a9233b36a13a758840522b78914 yt/data_objects/static_output.py
--- a/yt/data_objects/static_output.py
+++ b/yt/data_objects/static_output.py
@@ -106,8 +106,32 @@
             return self.__getattribute__(attr)
         return fnc
 
+    _field_types = None
+    @property
+    def field_types(self):
+        if self._field_types is None:
+            self._field_types = set(t for t, n in self.ds.field_info)
+        return self._field_types
+
     def __dir__(self):
-        return list(set(t for t, n in self.ds.field_info))
+        return list(self.field_types)
+
+    def __iter__(self):
+        for ft in self.field_types:
+            fnc = FieldNameContainer(self.ds, ft)
+            if len(dir(fnc)) == 0:
+                yield self.__getattribute__(ft)
+            else:
+                yield fnc
+
+    def __contains__(self, obj):
+        ob = None
+        if isinstance(obj, FieldNameContainer):
+            ob = obj.field_type
+        elif isinstance(obj, string_types):
+            ob = obj
+
+        return ob in self.field_types
 
 class FieldNameContainer(object):
     def __init__(self, ds, field_type):
@@ -125,6 +149,26 @@
         return [n for t, n in self.ds.field_info
                 if t == self.field_type]
 
+    def __iter__(self):
+        for t, n in self.ds.field_info:
+            if t == self.field_type:
+                yield self.ds.field_info[t, n]
+
+    def __contains__(self, obj):
+        if isinstance(obj, DerivedField):
+            if self.field_type == obj.name[0] and obj.name in self.ds.field_info:
+                # e.g. from a completely different dataset
+                if self.ds.field_info[obj.name] is not obj:
+                    return False
+                return True
+        elif isinstance(obj, tuple):
+            if self.field_type == obj[0] and obj in self.ds.field_info:
+                return True
+        elif isinstance(obj, string_types):
+            if (self.field_type, obj) in self.ds.field_info:
+                return True
+        return False
+
 class IndexProxy(object):
     # This is a simple proxy for Index objects.  It enables backwards
     # compatibility so that operations like .h.sphere, .h.print_stats and

diff -r bcdf7ee3c05bc4c96372fc3bfa5a581595c09bb2 -r c5bdb1f946cd9a9233b36a13a758840522b78914 yt/fields/tests/test_field_name_container.py
--- a/yt/fields/tests/test_field_name_container.py
+++ b/yt/fields/tests/test_field_name_container.py
@@ -1,13 +1,28 @@
-from yt.utilities.answer_testing.framework import \
-    requires_ds, \
-    data_dir_load
+from yt import \
+    load
+from yt.testing import \
+    requires_file
+
+def do_field_type(ft):
+    for field_name in dir(ft):
+        f = getattr(ft, field_name)
+        assert ((ft.field_type, field_name) == f.name)
+    for field in ft:
+        f = getattr(ft, field.name[1])
+        assert (f == field)
+        assert (f in ft)
+        assert (f.name in ft)
+        assert (f.name[1] in ft)
+
 
 enzotiny = "enzo_tiny_cosmology/DD0046/DD0046"
- at requires_ds(enzotiny)
-def test_simulated_halo_mass_function():
-    ds = data_dir_load(enzotiny)
+ at requires_file(enzotiny)
+def test_field_name_container():
+    ds = load(enzotiny)
     for field_type in dir(ds.fields):
+        assert (field_type in ds.fields)
         ft = getattr(ds.fields, field_type)
-        for field_name in dir(ft):
-            f = getattr(ft, field_name)
-            assert ((field_type, field_name) == f.name)
+        do_field_type(ft)
+    for field_type in ds.fields:
+        assert (field_type in ds.fields)
+        do_field_type(field_type)

Repository URL: https://bitbucket.org/yt_analysis/yt/

--

This is a commit notification from bitbucket.org. You are receiving
this because you have the service enabled, addressing the recipient of
this email.