streamlit-octostar-utils 0.5.0.dev1__tar.gz → 0.5.0.dev2__tar.gz

{streamlit_octostar_utils-0.5.0.dev1 → streamlit_octostar_utils-0.5.0.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: streamlit-octostar-utils
-Version: 0.5.0.dev1
+Version: 0.5.0.dev2
 Summary: 
 License: MIT
 License-File: LICENSE

{streamlit_octostar_utils-0.5.0.dev1 → streamlit_octostar_utils-0.5.0.dev2}/pyproject.toml

@@ -5,7 +5,7 @@ include = '\.pyi?$'
 
 [tool.poetry]
 name = "streamlit-octostar-utils"
-version = "0.5.0-dev.1"
+version = "0.5.0-dev.2"
 description = ""
 license = "MIT"
 authors = ["Octostar"]

{streamlit_octostar_utils-0.5.0.dev1 → streamlit_octostar_utils-0.5.0.dev2}/streamlit_octostar_utils/api_crafter/nifi.py

@@ -238,6 +238,156 @@ class NifiFragmenter(object):
             pointer.get("merge_params") or {}, defragmenter_config, lambda _, v2: v2
         )
 
+    _REQUIRED_FRAGMENT_FIELDS = ("index", "count", "identifier")
+
+    @staticmethod
+    def get_fragment_info(entity, fragmenter_keylist):
+        """Read fragment metadata (identifier, count, index, root_uid, merge_params)
+        for a given fragmenter level. Read-only -- does not mutate the entity.
+
+        Args:
+            entity: A NifiEntity or NifiEntityProxy.
+            fragmenter_keylist: Dot-separated key path into the fragment config
+                (e.g. "document_pages" or "audio_split").
+
+        Returns:
+            dict with keys like identifier, count, index, root_uid, merge_params.
+            Empty dict if fragmenter_keylist is empty or intermediate keys are
+            missing (entity not fragmented at this level).
+
+        Raises:
+            KeyError: If the final fragment key is missing from the config.
+            ValueError: If the fragment info exists but lacks required fields
+                (index, count, identifier).
+        """
+        if not fragmenter_keylist:
+            return {}
+        pointer = entity.request["config"]["fragment"]
+        for k in fragmenter_keylist.split(".")[:-1]:
+            if not pointer.get(k):
+                return {}
+            pointer = pointer[k]
+        info = pointer[fragmenter_keylist.split(".")[-1]]
+        missing = [f for f in NifiFragmenter._REQUIRED_FRAGMENT_FIELDS if f not in info]
+        if missing:
+            raise RuntimeError(
+                f"Fragment info for '{fragmenter_keylist}' is missing required "
+                f"field(s): {', '.join(missing)}"
+            )
+        return info
+
+    @staticmethod
+    def identify_fragment_groups(nifi_batches):
+        """Find all fragmented entities grouped by their active fragmenter level.
+
+        Args:
+            nifi_batches: List of NifiEntityBatch objects.
+
+        Returns:
+            dict mapping fragmenter_keylist to list of entities at that level.
+            Empty dict if no fragments found. Callers use get_fragment_info()
+            to fetch metadata per entity, and filter by index==0 to find roots.
+        """
+        all_entities = list(itertools.chain(*[b.entities for b in nifi_batches]))
+        groups = {}
+        for e in all_entities:
+            stack = e.request["config"].get("fragment", {}).get("fragments_stack", [])
+            if stack:
+                groups.setdefault(stack[0], []).append(e)
+        return groups
+
+    @staticmethod
+    def build_fragment_tree_from_children_entities(root_entity, fragmenter_keylist):
+        """Recursively build a tree from a root fragment entity by walking
+        its children_entities.
+
+        Args:
+            root_entity: The root entity (index 0) to start from.
+            fragmenter_keylist: The fragmenter level to build for.
+
+        Returns:
+            Nested dict with keys:
+                "entity": NifiEntity/NifiEntityProxy
+                "index": int
+                "merge_params": dict or None
+                "children": list of child trees
+        """
+        info = NifiFragmenter.get_fragment_info(root_entity, fragmenter_keylist)
+        child_fragments = []
+        for e in root_entity.children_entities:
+            try:
+                child_info = NifiFragmenter.get_fragment_info(e, fragmenter_keylist)
+                if child_info:
+                    child_fragments.append(e)
+            except (AttributeError, KeyError):
+                pass
+        return {
+            "entity": root_entity,
+            "index": info.get("index"),
+            "merge_params": info.get("merge_params"),
+            "children": [
+                NifiFragmenter.build_fragment_tree_from_children_entities(child, fragmenter_keylist)
+                for child in child_fragments
+            ],
+        }
+
+    @staticmethod
+    def extract_tree_entities(tree):
+        """Flatten a fragment tree into a list of all entities (pre-order).
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+
+        Returns:
+            List of entities in pre-order traversal.
+        """
+        entities = [tree["entity"]]
+        for child in tree.get("children", []):
+            entities.extend(NifiFragmenter.extract_tree_entities(child))
+        return entities
+
+    @staticmethod
+    def iterate_fragments_tree(tree, order="post"):
+        """Yield tree nodes in traversal order.
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+            order: "post" (children first, default) or "pre" (parent first).
+
+        Yields:
+            dict nodes with "entity", "index", "merge_params", "children" keys.
+        """
+        children = sorted(tree.get("children", []), key=lambda x: x["index"])
+        if order == "pre":
+            yield tree
+        for child in children:
+            yield from NifiFragmenter.iterate_fragments_tree(child, order)
+        if order == "post":
+            yield tree
+
+    @staticmethod
+    def reduce_fragments_tree(tree, leaf_fn, parent_fn):
+        """Bottom-up tree reduction. Processes leaves first, then folds results up.
+
+        Args:
+            tree: Fragment tree node (from build_fragment_tree_from_children_entities).
+            leaf_fn: Callable(node) -> result, called on nodes with no children.
+            parent_fn: Callable(node, child_results) -> result, called on
+                nodes with children. child_results is a list of results from
+                child nodes, sorted by index.
+
+        Returns:
+            The result from the root node.
+        """
+        children = sorted(tree.get("children", []), key=lambda x: x["index"])
+        if not children:
+            return leaf_fn(tree)
+        child_results = [
+            NifiFragmenter.reduce_fragments_tree(child, leaf_fn, parent_fn)
+            for child in children
+        ]
+        return parent_fn(tree, child_results)
+
 
 class NifiEntityBatch(object):
     def __init__(self, entities, config, config_key):
@@ -777,25 +927,68 @@ class NifiEntity(object):
             entity_type = self.record["entity_type"]
             return entity_type == type or type in self.request["ontology_info"]["parents"]
 
-    def is_fragmented(self) -> bool:
-        return bool(self.request["config"].get("fragment", {}).get("fragments_stack"))
-
-    def is_root_fragment(self, entity) -> bool:
-        def _is_sub_fragment_recursive(fragment: dict) -> bool:
-            if not isinstance(fragment, dict):
-                return False
-            if all(k in fragment for k in ["index", "count", "identifier"]):
-                return fragment.get("index", 0) != 0
-            for value in fragment.values():
-                if isinstance(value, dict):
-                    if _is_sub_fragment_recursive(value):
-                        return True
-            return False
-
+    def is_fragmented(self, fragment_name_or_idx=None) -> bool:
+        """Check whether this entity is part of a fragmentation.
+
+        Args:
+            fragment_name_or_idx: If None (default), returns True if the entity
+                belongs to any fragmentation level. If an int, checks whether the
+                fragments_stack has an entry at that index. If a string, checks
+                whether that fragmenter keylist is present in the stack.
+
+        Returns:
+            True if the entity is fragmented (at the specified level, if given).
+        """
+        stack = self.request["config"].get("fragment", {}).get("fragments_stack", [])
+        if fragment_name_or_idx is None:
+            return bool(stack)
+        if isinstance(fragment_name_or_idx, int):
+            return abs(fragment_name_or_idx) <= len(stack)
+        return fragment_name_or_idx in stack
+
+    def is_root_fragment(self, fragment_name_or_idx=-1, recurse=True) -> bool:
+        """Check whether this entity is a root fragment (index == 0).
+
+        Args:
+            fragment_name_or_idx: Which fragmentation level to check. An int
+                indexes into fragments_stack (default -1 = oldest level),
+                a string matches by fragmenter keylist name.
+            recurse: If True (default), check from the starting level towards
+                index 0 (most recent) and return True only if the entity is
+                root at every checked level. If False, only check the single
+                specified level.
+
+        Returns:
+            True if the entity is a root fragment at the specified level(s),
+            or if the entity is not fragmented at all.
+        """
         if not self.is_fragmented():
             return True
-        fragment = entity.request.get("config", {}).get("fragment", {})
-        return not _is_sub_fragment_recursive(fragment)
+        fragments_stack = self.request["config"]["fragment"]["fragments_stack"]
+
+        if isinstance(fragment_name_or_idx, int):
+            try:
+                resolved = (
+                    fragment_name_or_idx
+                    if fragment_name_or_idx >= 0
+                    else len(fragments_stack) + fragment_name_or_idx
+                )
+                key = fragments_stack[resolved]
+            except (IndexError, ValueError):
+                return True
+        else:
+            key = fragment_name_or_idx
+            if key not in fragments_stack:
+                return True
+            resolved = fragments_stack.index(key)
+
+        if recurse:
+            keys_to_check = fragments_stack[:resolved + 1]
+            return all(
+                NifiFragmenter.get_fragment_info(self, k).get("index", 0) == 0
+                for k in keys_to_check
+            )
+        return NifiFragmenter.get_fragment_info(self, key).get("index", 0) == 0
 
     def get_fragment_root_uid(self, fragment_name_or_idx) -> str:
         fragment_config = self.request.get("config", {}).get("fragment", {})