<commit_before>__version__ = "4.1.0.0"

import collections
import json
import inspect
from StringIO import StringIO
import bindings.maec_bundle as bundle_binding
import bindings.maec_package as package_binding
from cybox import Entity as cyboxEntity
from cybox.utils import Namespace
from maec.utils import maecMETA

def get_xmlns_string(ns_set):
    """Build a string with 'xmlns' definitions for every namespace in ns_set.

    Arguments:
    - ns_set: a set (or other iterable) of Namespace objects
    """
    xmlns_format = 'xmlns:{0.prefix}="{0.name}"'
    return "\n\t".join([xmlns_format.format(x) for x in ns_set])


def get_schemaloc_string(ns_set):
    """Build a "schemaLocation" string for every namespace in ns_set.

    Arguments:
    - ns_set: a set (or other iterable) of Namespace objects
    """
    schemaloc_format = '{0.name} {0.schema_location}'
    # Only include schemas that have a schema_location defined (for instance,
    # 'xsi' does not.
    return " ".join([schemaloc_format.format(x) for x in ns_set
                     if x.schema_location])

class Entity(object):
    """Base class for all classes in the MAEC SimpleAPI."""

    # By default (unless a particular subclass states otherwise), try to "cast"
    # invalid objects to the correct class using the constructor. Entity
    # subclasses should either provide a "sane" constructor or set this to
    # False.
    _try_cast = True

    def __init__(self):
        self._fields = {}

    @classmethod
    def _get_vars(cls):
        var_list = []
        for (name, obj) in inspect.getmembers(cls, inspect.isdatadescriptor):
            if isinstance(obj, TypedField):
                var_list.append(obj)

        return var_list

    def __eq__(self, other):
        # This fixes some strange behavior where an object isn't equal to
        # itself
        if other is self:
            return True

        # I'm not sure about this, if we want to compare exact classes or if
        # various subclasses will also do (I think not), but for now I'm going
        # to assume they must be equal. - GTB
        if self.__class__ != other.__class__:
            return False

        var_list = self.__class__._get_vars()

        # If there are no TypedFields, assume this class hasn't been
        # "TypedField"-ified, so we don't want these to inadvertently return
        # equal.
        if not var_list:
            return False

        for f in var_list:
            if not f.comparable:
                continue
            if getattr(self, f.attr_name) != getattr(other, f.attr_name):
                return False

        return True

    def __ne__(self, other):
        return not self == other

    def to_obj(self):
        """Default implementation of a to_obj function.

        Subclasses can override this function."""

        entity_obj = self._binding_class()

        for field in self.__class__._get_vars():
            val = getattr(self, field.attr_name)

            if field.multiple and val:
                val = [x.to_obj() for x in val]
            elif isinstance(val, Entity):
                val = val.to_obj()

            setattr(entity_obj, field.name, val)

        self._finalize_obj(entity_obj)

        return entity_obj

    def _finalize_obj(self, entity_obj):
        """Subclasses can define additional items in the binding object.

        `entity_obj` should be modified in place.
        """
        pass

    def to_dict(self):
        """Default implementation of a to_dict function.

        Subclasses can override this function."""

        entity_dict = {}

        for field in self.__class__._get_vars():
            val = getattr(self, field.attr_name)

            if field.multiple and val:
                val = [x.to_dict() for x in val]
            if isinstance(val, Entity):
                val = val.to_dict()

            # Only add non-None objects or non-empty lists
            if val is not None and val != []:
                entity_dict[field.key_name] = val

        self._finalize_dict(entity_dict)

        return entity_dict

    def _finalize_dict(self, entity_dict):
        """Subclasses can define additional items in the dictionary.

        `entity_dict` should be modified in place.
        """
        pass

    @classmethod
    def from_obj(cls, cls_obj=None):
        if not cls_obj:
            return None

        entity = cls()

        for field in cls._get_vars():
            val = getattr(cls_obj, field.name)
            if field.type_:
                if field.multiple and val is not None:
                    val = [field.type_.from_obj(x) for x in val]
                else:
                    val = field.type_.from_obj(val)
            setattr(entity, field.attr_name, val)

        return entity

    @classmethod
    def from_dict(cls, cls_dict=None):
        if cls_dict is None:
            return None

        entity = cls()

        # Shortcut if an actual dict is not provided:
        if not isinstance(cls_dict, dict):
            value = cls_dict
            # Call the class's constructor
            try:
                return cls(value)
            except TypeError:
                raise TypeError("Could not instantiate a %s from a %s: %s" %
                                (cls, type(value), value))

        for field in cls._get_vars():
            val = cls_dict.get(field.key_name)
            if field.type_:
                if issubclass(field.type_, EntityList):
                    val = field.type_.from_list(val)
                elif field.multiple:
                    if val is not None:
                        val = [field.type_.from_dict(x) for x in val]
                    else:
                        val = []
                else:
                    val = field.type_.from_dict(val)
            setattr(entity, field.attr_name, val)

        return entity

    def to_xml(self, include_namespaces=True, namespace_dict=None,
               pretty=True):
        """Export an object as an XML String.

        :param include_namespaces: whether to include xmlns and
          xsi:schemaLocation attributes on the root element. Set to true by
          default.
        :type include_namespaces: bool
        :param namespace_dict: mapping of additional XML namespaces to prefixes
        :type namespace_dict: dict
        :param pretty: produce readable (``True``) or compact (``False``)
          output. Default is ``True``
        :type pretty: bool
        """
        namespace_def = ""

        if include_namespaces:
            namespace_def = self._get_namespace_def(namespace_dict)

        if not pretty:
            namespace_def = namespace_def.replace('\n\t', ' ')

        s = StringIO()
        self.to_obj().export(s, 0, namespacedef_=namespace_def,
                             pretty_print=pretty)
        return s.getvalue()

    def to_xml_file(self, filename):
        """Export an object to an XML file. Only supports Package or Bundle objects at the moment."""
        out_file  = open(filename, 'w')
        out_file.write("<?xml version='1.0' encoding='UTF-8'?>\n")
        self.to_obj().export(out_file, 0, namespace_def = self._get_namespace_def(namespace_dict))
        out_file.close()

    def to_json(self):
        """Export an object as a JSON string.
        """
        return json.dumps(self.to_dict())

    def _get_namespace_def(self, additional_ns_dict=None):
        # copy necessary namespaces

        namespaces = self._get_namespaces()

        if additional_ns_dict:
            for ns, prefix in additional_ns_dict.iteritems():
                namespaces.update([Namespace(ns, prefix)])

        # if there are any other namepaces, include xsi for "schemaLocation"
        if namespaces:
            namespaces.update([maecMETA.lookup_prefix('xsi')])
<commit_msg>Added optional namespace_dict parameter to to_xml_file()<commit_after>__version__ = "4.1.0.0"

import collections
import json
import inspect
from StringIO import StringIO
import bindings.maec_bundle as bundle_binding
import bindings.maec_package as package_binding
from cybox import Entity as cyboxEntity
from cybox.utils import Namespace
from maec.utils import maecMETA

def get_xmlns_string(ns_set):
    """Build a string with 'xmlns' definitions for every namespace in ns_set.

    Arguments:
    - ns_set: a set (or other iterable) of Namespace objects
    """
    xmlns_format = 'xmlns:{0.prefix}="{0.name}"'
    return "\n\t".join([xmlns_format.format(x) for x in ns_set])


def get_schemaloc_string(ns_set):
    """Build a "schemaLocation" string for every namespace in ns_set.

    Arguments:
    - ns_set: a set (or other iterable) of Namespace objects
    """
    schemaloc_format = '{0.name} {0.schema_location}'
    # Only include schemas that have a schema_location defined (for instance,
    # 'xsi' does not.
    return " ".join([schemaloc_format.format(x) for x in ns_set
                     if x.schema_location])

class Entity(object):
    """Base class for all classes in the MAEC SimpleAPI."""

    # By default (unless a particular subclass states otherwise), try to "cast"
    # invalid objects to the correct class using the constructor. Entity
    # subclasses should either provide a "sane" constructor or set this to
    # False.
    _try_cast = True

    def __init__(self):
        self._fields = {}

    @classmethod
    def _get_vars(cls):
        var_list = []
        for (name, obj) in inspect.getmembers(cls, inspect.isdatadescriptor):
            if isinstance(obj, TypedField):
                var_list.append(obj)

        return var_list

    def __eq__(self, other):
        # This fixes some strange behavior where an object isn't equal to
        # itself
        if other is self:
            return True

        # I'm not sure about this, if we want to compare exact classes or if
        # various subclasses will also do (I think not), but for now I'm going
        # to assume they must be equal. - GTB
        if self.__class__ != other.__class__:
            return False

        var_list = self.__class__._get_vars()

        # If there are no TypedFields, assume this class hasn't been
        # "TypedField"-ified, so we don't want these to inadvertently return
        # equal.
        if not var_list:
            return False

        for f in var_list:
            if not f.comparable:
                continue
            if getattr(self, f.attr_name) != getattr(other, f.attr_name):
                return False

        return True

    def __ne__(self, other):
        return not self == other

    def to_obj(self):
        """Default implementation of a to_obj function.

        Subclasses can override this function."""

        entity_obj = self._binding_class()

        for field in self.__class__._get_vars():
            val = getattr(self, field.attr_name)

            if field.multiple and val:
                val = [x.to_obj() for x in val]
            elif isinstance(val, Entity):
                val = val.to_obj()

            setattr(entity_obj, field.name, val)

        self._finalize_obj(entity_obj)

        return entity_obj

    def _finalize_obj(self, entity_obj):
        """Subclasses can define additional items in the binding object.

        `entity_obj` should be modified in place.
        """
        pass

    def to_dict(self):
        """Default implementation of a to_dict function.

        Subclasses can override this function."""

        entity_dict = {}

        for field in self.__class__._get_vars():
            val = getattr(self, field.attr_name)

            if field.multiple and val:
                val = [x.to_dict() for x in val]
            if isinstance(val, Entity):
                val = val.to_dict()

            # Only add non-None objects or non-empty lists
            if val is not None and val != []:
                entity_dict[field.key_name] = val

        self._finalize_dict(entity_dict)

        return entity_dict

    def _finalize_dict(self, entity_dict):
        """Subclasses can define additional items in the dictionary.

        `entity_dict` should be modified in place.
        """
        pass

    @classmethod
    def from_obj(cls, cls_obj=None):
        if not cls_obj:
            return None

        entity = cls()

        for field in cls._get_vars():
            val = getattr(cls_obj, field.name)
            if field.type_:
                if field.multiple and val is not None:
                    val = [field.type_.from_obj(x) for x in val]
                else:
                    val = field.type_.from_obj(val)
            setattr(entity, field.attr_name, val)

        return entity

    @classmethod
    def from_dict(cls, cls_dict=None):
        if cls_dict is None:
            return None

        entity = cls()

        # Shortcut if an actual dict is not provided:
        if not isinstance(cls_dict, dict):
            value = cls_dict
            # Call the class's constructor
            try:
                return cls(value)
            except TypeError:
                raise TypeError("Could not instantiate a %s from a %s: %s" %
                                (cls, type(value), value))

        for field in cls._get_vars():
            val = cls_dict.get(field.key_name)
            if field.type_:
                if issubclass(field.type_, EntityList):
                    val = field.type_.from_list(val)
                elif field.multiple:
                    if val is not None:
                        val = [field.type_.from_dict(x) for x in val]
                    else:
                        val = []
                else:
                    val = field.type_.from_dict(val)
            setattr(entity, field.attr_name, val)

        return entity

    def to_xml(self, include_namespaces=True, namespace_dict=None,
               pretty=True):
        """Export an object as an XML String.

        :param include_namespaces: whether to include xmlns and
          xsi:schemaLocation attributes on the root element. Set to true by
          default.
        :type include_namespaces: bool
        :param namespace_dict: mapping of additional XML namespaces to prefixes
        :type namespace_dict: dict
        :param pretty: produce readable (``True``) or compact (``False``)
          output. Default is ``True``
        :type pretty: bool
        """
        namespace_def = ""

        if include_namespaces:
            namespace_def = self._get_namespace_def(namespace_dict)

        if not pretty:
            namespace_def = namespace_def.replace('\n\t', ' ')

        s = StringIO()
        self.to_obj().export(s, 0, namespacedef_=namespace_def,
                             pretty_print=pretty)
        return s.getvalue()

    def to_xml_file(self, filename, namespace_dict=None):
        """Export an object to an XML file. Only supports Package or Bundle objects at the moment."""
        out_file  = open(filename, 'w')
        out_file.write("<?xml version='1.0' encoding='UTF-8'?>\n")
        self.to_obj().export(out_file, 0, namespace_def = self._get_namespace_def(namespace_dict))
        out_file.close()

    def to_json(self):
        """Export an object as a JSON string.
        """
        return json.dumps(self.to_dict())

    def _get_namespace_def(self, additional_ns_dict=None):
        # copy necessary namespaces

        namespaces = self._get_namespaces()

        if additional_ns_dict:
            for ns, prefix in additional_ns_dict.iteritems():
                namespaces.update([Namespace(ns, prefix)])

        # if there are any other namepaces, include xsi for "schemaLocation"
        if namespaces:
            namespaces.update([maecMETA.lookup_prefix('xsi')])

{'fraction_non_alphanumeric': 0.08430675469781615, 'fraction_numerical': 0.0016505840528186896, 'mean_word_length': 1.9931597947938438, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 2, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 19}, 'pii_count': 2, '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': '1825579', 'n_tokens_mistral': 4459, 'n_tokens_neox': 4266, 'n_words': 1396}