<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: <NAME> <<EMAIL>>

"""
This module contains code to convert discourse graphs into rooted, ordered
trees.
"""

from collections import defaultdict, deque
import textwrap

from nltk.tree import ParentedTree

from discoursegraphs import (
    EdgeTypes, istoken, select_neighbors_by_edge_attribute)


class DGParentedTree(ParentedTree):
    """An nltk.tree.ParentedTree with an additional root_id parameter."""
    def __init__(self, node, children=None, root_id=None):
        # super calls __init__() of base class nltk.tree.ParentedTree
        super(DGParentedTree, self).__init__(node, children)
        self.root_id = root_id

    def get_position(self, rst_tree, node_id=None):
        """Get the linear position of an element of this DGParentedTree in a RSTTree.

        If ``node_id`` is given, this will return the position of the subtree
        with that node ID. Otherwise, the position of the root of this
        DGParentedTree in the given RSTTree.
        """
        if node_id is None:
            node_id = self.root_id

        if node_id in rst_tree.edu_set:
            return rst_tree.edus.index(node_id)

        return min(self.get_position(rst_tree, child_node_id)
                   for child_node_id in rst_tree.child_dict[node_id])

    def node_height(self, rst_tree, node_id=None):
        if node_id is None:
            node_id = self.root_id
        return rst_tree.node_height(node_id)


def debug_root_label(root_label, debug=False, root_id=None):
    if debug is True and root_id is not None:
        return root_label + " ({})".format(root_id)
    else:
        return root_label


def t(root, children=None, debug=False, root_id=None):
    "Create DGParentedTree from a root (str) and a list of (str, list) tuples."
    if isinstance(root, DGParentedTree):
        if children is None:
            return root
        return DGParentedTree(root, children, root_id)

    elif isinstance(root, basestring):
        root = debug_root_label(root, debug, root_id)

        # Beware: DGParentedTree is a subclass of list!
        if isinstance(children, DGParentedTree):
            child_trees = [children]

        elif isinstance(children, list):
            child_trees = []
            for child in children:
                if isinstance(child, DGParentedTree):
                    child_trees.append(child)
                elif isinstance(child, list):
                    child_trees.extend(child)
                elif isinstance(child, tuple):
                    child_trees.append(t(*child))
                elif isinstance(child, basestring):
                    child_trees.append(child)
                else:
                    raise NotImplementedError

        elif isinstance(children, basestring):
            # this tree does only have one child, a leaf node
            # TODO: this is a workaround for the following problem:
            # DGParentedTree('foo', [DGParentedTree('bar', [])]) != DGParentedTree('foo', ['bar'])
            child_trees = [DGParentedTree(children, [])]

        else:
            # this tree only consists of one leaf node
            assert children is None
            child_trees = []

        return DGParentedTree(root, child_trees, root_id)

    else:
        raise NotImplementedError




def get_child_nodes(docgraph, parent_node_id, data=False):
    """Yield all nodes that the given node dominates or spans."""
    return select_neighbors_by_edge_attribute(
        docgraph=docgraph,
        source=parent_node_id,
        attribute='edge_type',
        value=[EdgeTypes.dominance_relation],
        data=data)


def get_parents(docgraph, child_node, strict=True):
    """Return a list of parent nodes that dominate this child.

    In a 'syntax tree' a node never has more than one parent node
    dominating it. To enforce this, set strict=True.
<commit_msg>simplified t() to work with both Tree and DGParentedTree<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: <NAME> <<EMAIL>>

"""
This module contains code to convert discourse graphs into rooted, ordered
trees.
"""

from collections import defaultdict, deque
import textwrap

from nltk.tree import Tree, ParentedTree

from discoursegraphs import (
    EdgeTypes, istoken, select_neighbors_by_edge_attribute)


class DGParentedTree(ParentedTree):
    """An nltk.tree.ParentedTree with an additional root_id parameter."""
    def __init__(self, node, children=None, root_id=None):
        # super calls __init__() of base class nltk.tree.ParentedTree
        super(DGParentedTree, self).__init__(node, children)
        self.root_id = root_id

    def get_position(self, rst_tree, node_id=None):
        """Get the linear position of an element of this DGParentedTree in an RSTTree.

        If ``node_id`` is given, this will return the position of the subtree
        with that node ID. Otherwise, the position of the root of this
        DGParentedTree in the given RSTTree is returned.
        """
        if node_id is None:
            node_id = self.root_id

        if node_id in rst_tree.edu_set:
            return rst_tree.edus.index(node_id)

        return min(self.get_position(rst_tree, child_node_id)
                   for child_node_id in rst_tree.child_dict[node_id])

    def node_height(self, rst_tree, node_id=None):
        if node_id is None:
            node_id = self.root_id
        return rst_tree.node_height(node_id)


def debug_root_label(root_label, debug=False, root_id=None):
    if debug is True and root_id is not None:
        return root_label + " ({})".format(root_id)
    else:
        return root_label


def t(root, children=None, debug=False, root_id=None):
    "Create (DGParented)Tree from a root (str) and a list of (str, list) tuples."
    if isinstance(root, Tree):
        if children is None:
            return root
        return root.__class__(root, children, root_id)

    elif isinstance(root, basestring):
        root = debug_root_label(root, debug, root_id)

        # Beware: (DGParented)Tree is a subclass of list!
        if isinstance(children, Tree):
            child_trees = [children]

        elif isinstance(children, list):
            child_trees = []
            for child in children:
                if isinstance(child, Tree):
                    child_trees.append(child)
                elif isinstance(child, list):
                    child_trees.extend(child)
                elif isinstance(child, tuple):
                    child_trees.append(t(*child))
                elif isinstance(child, basestring):
                    child_trees.append(child)
                else:
                    raise NotImplementedError

        elif isinstance(children, basestring):
            # this tree does only have one child, a leaf node
            # TODO: this is a workaround for the following problem:
            # Tree('foo', [Tree('bar', [])]) != Tree('foo', ['bar'])
            child_trees = [Tree(children, [])]

        else:
            # this tree only consists of one leaf node
            assert children is None
            child_trees = []

        return DGParentedTree(root, child_trees, root_id)

    else:
        raise NotImplementedError




def get_child_nodes(docgraph, parent_node_id, data=False):
    """Yield all nodes that the given node dominates or spans."""
    return select_neighbors_by_edge_attribute(
        docgraph=docgraph,
        source=parent_node_id,
        attribute='edge_type',
        value=[EdgeTypes.dominance_relation],
        data=data)


def get_parents(docgraph, child_node, strict=True):
    """Return a list of parent nodes that dominate this child.

    In a 'syntax tree' a node never has more than one parent node
    dominating it. To enforce this, set strict=True.

{'fraction_non_alphanumeric': 0.08416220351951033, 'fraction_numerical': 0.0002550369803621525, 'mean_word_length': 2.521778176919623, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2349565', 'n_tokens_mistral': 2330, 'n_tokens_neox': 2220, 'n_words': 689}