"""Classes to represent and query categorical systems."""
import datetime
import functools
import importlib
import importlib.resources
import itertools
import pathlib
import pickle
import typing
from typing import TypeVar
import natsort
import networkx as nx
import pandas
import strictyaml as sy
from black import Mode, format_str
from ruamel.yaml import YAML
from . import data
from ._conversions import Conversion, ConversionSpec
# Categorization, or any subclass.
CategorizationT = TypeVar("CategorizationT", bound="Categorization")
[docs]
class Category:
"""A single category."""
_strictyaml_schema = sy.Map(
{
"title": sy.Str(),
sy.Optional("comment"): sy.Str(),
sy.Optional("alternative_codes"): sy.Seq(sy.Str()),
sy.Optional("info"): sy.MapPattern(sy.Str(), sy.Any()),
}
)
def __init__(
self,
codes: tuple[str, ...],
categorization: "Categorization",
title: str,
comment: None | str = None,
info: None | dict = None,
):
self.codes = codes
self.title = title
self.comment = comment
self.categorization = categorization
if info is None:
self.info = {}
else:
self.info = info
self._hash = None
@classmethod
def from_spec(cls, code: str, spec: dict, categorization: "Categorization"):
codes = [code]
if "alternative_codes" in spec:
codes += spec["alternative_codes"]
del spec["alternative_codes"]
return cls(
codes=tuple(codes),
categorization=categorization,
title=spec["title"],
comment=spec.get("comment"),
info=spec.get("info"),
)
[docs]
def to_spec(self) -> tuple[str, dict[str, str | dict | list]]:
"""Turn this category into a specification ready to be written to a yaml file.
Returns
-------
(code: str, spec: dict)
Primary code and specification dict
"""
code = self.codes[0]
spec: dict[str, str | dict | list[str]] = {"title": self.title}
if self.comment is not None:
spec["comment"] = self.comment
if len(self.codes) > 1:
spec["alternative_codes"] = list(self.codes[1:])
if self.info:
spec["info"] = self.info
return code, spec
def __str__(self) -> str:
return f"{self.codes[0]} {self.title}"
def __eq__(self, other: object):
if not isinstance(other, Category):
return NotImplemented
return any(x in other.codes for x in self.codes) and (
self.categorization is other.categorization
or self.categorization.name.startswith(f"{other.categorization.name}_")
or other.categorization.name.startswith(f"{self.categorization.name}_")
or self.categorization.name == other.categorization.name
)
def __repr__(self) -> str:
return f"<{self.categorization.name}: {self.codes[0]!r}>"
def __hash__(self):
if self._hash is None:
self._hash = hash(self.categorization.name + self.codes[0])
return self._hash
def __lt__(self, other):
s = natsort.natsorted((self.codes[0], other.codes[0]))
return s[0] == self.codes[0] and self != other
[docs]
class HierarchicalCategory(Category):
"""A single category from a HierarchicalCategorization."""
_strictyaml_schema = sy.Map(
{
"title": sy.Str(),
sy.Optional("comment"): sy.Str(),
sy.Optional("alternative_codes"): sy.Seq(sy.Str()),
sy.Optional("info"): sy.MapPattern(sy.Str(), sy.Any()),
sy.Optional("children"): sy.Seq(sy.Seq(sy.Str())),
}
)
def __init__(
self,
codes: tuple[str],
categorization: "HierarchicalCategorization",
title: str,
comment: None | str = None,
info: None | dict = None,
):
Category.__init__(self, codes, categorization, title, comment, info)
self.categorization = categorization
[docs]
def to_spec(self) -> tuple[str, dict[str, str | dict | list]]:
"""Turn this category into a specification ready to be written to a yaml file.
Returns
-------
(code: str, spec: dict)
Primary code and specification dict
"""
code, spec = Category.to_spec(self)
children = [
list(sorted(c.codes[0] for c in child_set)) for child_set in self.children
]
if children:
spec["children"] = children
return code, spec
@property
def children(self) -> list[set["HierarchicalCategory"]]:
"""The sets of subcategories comprising this category.
The first set is canonical, the other sets are alternative.
Only the canonical sets are used to calculate the level of a category."""
return self.categorization.children(self)
@property
def parents(self) -> set["HierarchicalCategory"]:
"""The super-categories where this category is a member of any set of children.
Note that all possible parents are returned, not "canonical" parents.
"""
return self.categorization.parents(self)
@property
def ancestors(self) -> set["HierarchicalCategory"]:
"""The super-categories where this category or any of its parents is a member
of any set of children, transitively.
Note that all possible ancestors are returned, not only "canonical" ones.
"""
return self.categorization.ancestors(self)
@property
def descendants(self) -> set["HierarchicalCategory"]:
"""The sets of subcategories comprising this category directly or indirectly.
Note that all possible descendants are returned, not only "canonical" ones."""
return self.categorization.descendants(self)
@property
def is_leaf(self) -> bool:
"""Is this category a leaf category, i.e. without children?"""
return not any(self.children)
@property
def leaf_children(self) -> list[set["HierarchicalCategory"]]:
"""The sets of subcategories which are descendants of this category and do not
have children themselves.
Sets of children are chased separately, so each set of leaf children is
self-sufficient to reconstruct this category (if the categorization allows
reconstructing categories from their children, i.e. if total_sum is set)."""
ret = []
for children in self.children:
n = []
for child in children:
if child.is_leaf:
n.append([{child}])
else:
n.append(child.leaf_children)
ret += [set(itertools.chain(*x)) for x in itertools.product(*n)]
return ret
@property
def level(self) -> int:
"""The level of the category.
The canonical top-level category has level 1 and its children have level 2 etc.
To calculate the level, only the first ("canonical") set of children is
considered for intermediate categories.
"""
return self.categorization.level(self)
[docs]
class Categorization:
"""A single categorization system.
A categorization system comprises a set of categories, and their relationships as
well as metadata describing the categorization system itself.
Use the categorization object like a dictionary, where codes can be translated
to their meaning using ``cat[code]`` and all codes are available using
``cat.keys()``. Metadata about the categorization is provided in attributes.
If `pandas` is available, you can access a `pandas.DataFrame` with all
category codes, and their meanings at ``cat.df``.
Attributes
----------
name : str
The unique name/code
references : str
Citable reference(s)
title : str
A short, descriptive title for humans
comment : str
Notes and explanations for humans
institution : str
Where the categorization originates
last_update : datetime.date
The date of the last change
version : str, optional
The version of the Categorization, if there are multiple versions
hierarchical : bool
True if descendants and ancestors are defined
"""
hierarchical: bool = False
_strictyaml_schema = sy.Map(
{
"name": sy.Str(),
"title": sy.Str(),
"comment": sy.Str(),
"references": sy.Str(),
"institution": sy.Str(),
"last_update": sy.Str(),
"hierarchical": sy.Bool(),
sy.Optional("version"): sy.Str(),
"categories": sy.MapPattern(sy.Str(), Category._strictyaml_schema),
}
)
def _add_categories(self, categories: dict[str, dict]):
for code, spec in categories.items():
cat = Category.from_spec(code=code, spec=spec, categorization=self)
self._primary_code_map[code] = cat
for icode in cat.codes:
self._all_codes_map[icode] = cat
def __init__(
self,
*,
categories: dict[str, dict],
name: str,
title: str,
comment: str,
references: str,
institution: str,
last_update: datetime.date,
version: None | str = None,
):
self._primary_code_map: dict[str, Category] = {}
self._all_codes_map: dict[str, Category] = {}
self.name = name
self.references = references
self.title = title
self.comment = comment
self.institution = institution
self.last_update = last_update
self.version = version
self._add_categories(categories)
# is filled in __init__.py to contain all categorizations
self._cats: dict[str, Categorization] = {}
def __hash__(self):
return hash(self.name)
[docs]
@classmethod
def from_yaml(
cls: type[CategorizationT],
filepath: str | pathlib.Path | typing.TextIO,
) -> CategorizationT:
"""Read Categorization from a StrictYaml file."""
try:
yaml = sy.load(filepath.read(), schema=cls._strictyaml_schema)
except AttributeError:
with open(filepath) as fd:
yaml = sy.load(fd.read(), schema=cls._strictyaml_schema)
return cls.from_spec(yaml.data)
[docs]
@classmethod
def from_spec(
cls: type[CategorizationT], spec: dict[str, typing.Any]
) -> CategorizationT:
"""Create Categorization from a Dictionary specification."""
if spec["hierarchical"] != cls.hierarchical:
raise ValueError(
"Specification is for a hierarchical categorization, use"
"HierarchicalCategorization.from_spec."
)
last_update = datetime.date.fromisoformat(spec["last_update"])
return cls(
categories=spec["categories"],
name=spec["name"],
title=spec["title"],
comment=spec["comment"],
references=spec["references"],
institution=spec["institution"],
last_update=last_update,
version=spec.get("version"),
)
[docs]
@staticmethod
def from_pickle(
filepath: str | pathlib.Path | typing.IO[bytes],
) -> CategorizationT:
"""De-serialize Categorization from a file written by to_pickle.
Note that this uses the pickle module, which executes arbitrary code in the
provided file. Only load from pickle files that you trust."""
return from_pickle(filepath)
[docs]
@staticmethod
def from_python(
filepath: str | pathlib.Path | typing.IO[bytes],
) -> CategorizationT:
"""De-serialize Categorization from a file written by to_python.
Note that this executes the python cache file. Only load from python cache files
you trust."""
return from_python(filepath)
[docs]
def to_spec(self) -> dict[str, typing.Any]:
"""Turn this categorization into a specification dictionary ready to be written
to a yaml file.
Returns
-------
spec: dict
Specification dictionary understood by `from_spec`.
"""
spec = {
"name": self.name,
"title": self.title,
"comment": self.comment,
"references": self.references,
"institution": self.institution,
"hierarchical": self.hierarchical,
"last_update": self.last_update.isoformat(),
}
if self.version is not None:
spec["version"] = self.version
categories = {}
for cat in self.values():
code, cat_spec = cat.to_spec()
categories[code] = cat_spec
spec["categories"] = categories
return spec
[docs]
def to_yaml(self, filepath: str | pathlib.Path) -> None:
"""Write to a YAML file."""
spec = self.to_spec()
yaml = YAML()
yaml.default_flow_style = False
with open(filepath, "w") as fd:
yaml.dump(spec, fd)
[docs]
def to_python(self, filepath: str | pathlib.Path) -> None:
"""Write spec to a Python file."""
spec = self.to_spec()
comment = (
"# Do not edit this file. It was auto-generated from the\n"
"# corresponding YAML file.\n"
)
with open(filepath, "w") as f:
f.write(comment)
f.write(f"spec = {format_str(repr(spec), mode=Mode())}")
[docs]
def to_pickle(self, filepath: str | pathlib.Path) -> None:
"""Serialize to a file using python's pickle."""
spec = self.to_spec()
with open(filepath, "wb") as fd:
pickle.dump(spec, fd, protocol=4)
[docs]
def keys(self) -> typing.KeysView[str]:
"""Iterate over the codes for all categories."""
return self._primary_code_map.keys()
[docs]
def values(self) -> typing.ValuesView[Category]:
"""Iterate over the categories."""
return self._primary_code_map.values()
[docs]
def items(self) -> typing.ItemsView[str, Category]:
"""Iterate over (primary code, category) pairs."""
return self._primary_code_map.items()
[docs]
def all_keys(self) -> typing.KeysView[str]:
"""Iterate over all codes for all categories."""
return self._all_codes_map.keys()
def __iter__(self) -> typing.Iterable[str]:
return iter(self._primary_code_map)
def __getitem__(self, code: str) -> Category:
"""Get the category for a code."""
return self._all_codes_map[code]
def __contains__(self, code: str) -> bool:
"""Can the code be mapped to a category?"""
return code in self._all_codes_map
def __len__(self) -> int:
return len(self._primary_code_map)
def __repr__(self) -> str:
return (
f"<Categorization {self.name} {self.title!r} with {len(self)} categories>"
)
def __str__(self) -> str:
return self.name
@property
def df(self) -> "pandas.DataFrame":
"""All category codes as a pandas dataframe."""
titles = []
comments = []
alternative_codes = []
for cat in self.values():
titles.append(cat.title)
comments.append(cat.comment)
alternative_codes.append(cat.codes[1:])
return pandas.DataFrame(
index=list(self.keys()),
data={
"title": titles,
"comment": comments,
"alternative_codes": alternative_codes,
},
)
def _extend_prepare(
self,
*,
categories: None | dict[str, dict] = None,
alternative_codes: None | dict[str, str] = None,
name: str,
title: None | str = None,
comment: None | str = None,
last_update: None | datetime.date = None,
) -> dict[str, typing.Any]:
spec = self.to_spec()
spec["name"] = f"{self.name}_{name}"
spec["references"] = ""
spec["institution"] = ""
if title is None:
spec["title"] = f"{self.title} + {name}"
else:
spec["title"] = self.title + title
if comment is None:
spec["comment"] = f"{self.comment} extended by {name}"
else:
spec["comment"] = self.comment + comment
if last_update is None:
spec["last_update"] = datetime.date.today().isoformat()
else:
spec["last_update"] = last_update.isoformat()
if categories is not None:
spec["categories"].update(categories)
if alternative_codes is not None:
for alias, primary in alternative_codes.items():
if "alternative_codes" not in spec["categories"][primary]:
spec["categories"][primary]["alternative_codes"] = []
spec["categories"][primary]["alternative_codes"].append(alias)
return spec
[docs]
def extend(
self: CategorizationT,
*,
categories: None | dict[str, dict] = None,
alternative_codes: None | dict[str, str] = None,
name: str,
title: None | str = None,
comment: None | str = None,
last_update: None | datetime.date = None,
) -> CategorizationT:
"""Extend the categorization with additional categories, yielding a new
categorization.
Metadata: the ``name``, ``title``, ``comment``, and ``last_update`` are updated
automatically (see below), the ``institution`` and ``references`` are deleted
and the values for ``version`` and ``hierarchical`` are kept.
You can set more accurate metadata (for example, your institution) on the
returned object if needed.
Parameters
----------
categories: dict, optional
Map of new category codes to their specification. The specification is a
dictionary with the keys "title", optionally "comment", and optionally
"alternative_codes".
alternative_codes: dict, optional
Map of new alternative codes. A dictionary with the new alternative code
as key and existing code as value.
name : str
The name of your extension. The returned Categorization will have a name
of "{old_name}_{name}", indicating that it is an extension of the underlying
Categorization.
title : str, optional
A string to add to the original title. If not provided, " + {name}" will be
used.
comment : str, optional
A string to add to the original comment. If not provided,
" extended by {name}" will be used.
last_update : datetime.date, optional
The date of the last update to this extension. Today will be used if not
provided.
Returns
-------
Extended categorization : Categorization
"""
spec = self._extend_prepare(
name=name,
categories=categories,
title=title,
comment=comment,
last_update=last_update,
alternative_codes=alternative_codes,
)
return Categorization.from_spec(spec)
def limit(
self: CategorizationT, categories: tuple[str, ...], *, name: str
) -> CategorizationT:
spec = self.to_spec()
spec["name"] = f"{self.name}_{name}"
spec["references"] = ""
spec["institution"] = ""
spec["title"] = f"{self.title} limited ({name})"
spec["comment"] = f"{self.comment} limited ({name})"
spec["last_update"] = datetime.date.today().isoformat()
spec["categories"] = {
code: cat for code, cat in spec["categories"].items() if code in categories
}
if isinstance(self, HierarchicalCategorization):
for cat in spec["categories"].values():
if "children" in cat:
new_children = []
for children in cat["children"]:
if all(x in categories for x in children):
new_children.append(children)
cat["children"] = new_children
return self.__class__.from_spec(spec)
def __eq__(self, other):
if not isinstance(other, Categorization):
return False
if self.name != other.name:
return False
return self._primary_code_map == other._primary_code_map
[docs]
def conversion_to(self, other: typing.Union["Categorization", str]) -> Conversion:
"""Get conversion to other categorization.
If conversion rules for this conversion are not included, raises
NotImplementedError."""
if isinstance(other, str):
other_name = other
else:
other_name = other.name
data_files = importlib.resources.files(data)
forward_file = data_files / f"conversion.{self.name}.{other_name}.csv"
if forward_file.is_file():
return ConversionSpec.from_csv(forward_file.open()).hydrate(cats=self._cats)
reverse_file = data_files / f"conversion.{other_name}.{self.name}.csv"
if reverse_file.is_file():
return (
ConversionSpec.from_csv(reverse_file.open())
.hydrate(cats=self._cats)
.reversed()
)
raise NotImplementedError(
f"Conversion between {self.name} and {other_name} not yet included."
)
[docs]
class HierarchicalCategorization(Categorization):
"""In a hierarchical categorization, descendants and ancestors (parents and
children) are defined for each category.
Attributes
----------
total_sum : bool
If the sum of the values of children equals the value of the parent for
extensive quantities. For example, a Categorization containing the Countries in
the EU and the EU could set `total_sum = True`, because the emissions of all
parts of the EU must equal the emissions of the EU. On the contrary, a
categorization of Industries with categories `Power:Fossil Fuels` and
`Power:Gas` which are both children of `Power` must set `total_sum = False`
to avoid double counting of fossil gas.
canonical_top_level_category : HierarchicalCategory
The level of a category is calculated with respect to the canonical top level
category. Commonly, this will be the world total or a similar category. If the
canonical top level category is not set (i.e. is ``None``), levels are not
defined for categories.
"""
hierarchical = True
_strictyaml_schema = sy.Map(
{
"name": sy.Str(),
"title": sy.Str(),
"comment": sy.Str(),
"references": sy.Str(),
"institution": sy.Str(),
"last_update": sy.Str(),
"hierarchical": sy.Bool(),
sy.Optional("version"): sy.Str(),
"total_sum": sy.Bool(),
sy.Optional("canonical_top_level_category"): sy.Str(),
"categories": sy.MapPattern(
sy.Str(), HierarchicalCategory._strictyaml_schema
),
}
)
def _add_categories(self, categories: dict[str, dict]):
for code, spec in categories.items():
cat = HierarchicalCategory.from_spec(
code=code, spec=spec, categorization=self
)
self._primary_code_map[code] = cat
self._graph.add_node(cat)
for icode in cat.codes:
self._all_codes_map[icode] = cat
for code, spec in categories.items():
if "children" in spec:
parent = self._all_codes_map[code]
for i, child_set in enumerate(spec["children"]):
for child_code in child_set:
self._graph.add_edge(
parent, self._all_codes_map[child_code], set=i
)
def __init__(
self,
*,
categories: dict[str, dict],
name: str,
title: str,
comment: str,
references: str,
institution: str,
last_update: datetime.date,
version: None | str = None,
total_sum: bool,
canonical_top_level_category: None | str = None,
):
self._graph = nx.MultiDiGraph()
Categorization.__init__(
self,
categories=categories,
name=name,
title=title,
comment=comment,
references=references,
institution=institution,
last_update=last_update,
version=version,
)
self.total_sum = total_sum
if canonical_top_level_category is None:
self.canonical_top_level_category: None | HierarchicalCategory = None
else:
self.canonical_top_level_category = self._all_codes_map[
canonical_top_level_category
]
def __getitem__(self, code: str) -> HierarchicalCategory:
"""Get the category for a code."""
return self._all_codes_map[code]
[docs]
def values(self) -> typing.ValuesView[HierarchicalCategory]:
"""Iterate over the categories."""
return self._primary_code_map.values()
[docs]
def items(self) -> typing.ItemsView[str, HierarchicalCategory]:
"""Iterate over (primary code, category) pairs."""
return self._primary_code_map.items()
[docs]
@classmethod
def from_spec(
cls: type[CategorizationT], spec: dict[str, typing.Any]
) -> CategorizationT:
"""Create Categorization from a Dictionary specification."""
if spec["hierarchical"] != cls.hierarchical:
raise ValueError(
"Specification is for a non-hierarchical categorization, use "
"Categorization.from_spec."
)
last_update = datetime.date.fromisoformat(spec["last_update"])
return cls(
categories=spec["categories"],
name=spec["name"],
title=spec["title"],
comment=spec["comment"],
references=spec["references"],
institution=spec["institution"],
last_update=last_update,
version=spec.get("version"),
total_sum=spec["total_sum"],
canonical_top_level_category=spec.get("canonical_top_level_category"),
)
[docs]
def to_spec(self) -> dict[str, typing.Any]:
"""Turn this categorization into a specification dictionary ready to be written
to a yaml file.
Returns
-------
spec: dict
Specification dictionary understood by `from_spec`.
"""
# we can't call Categorization.to_spec here because we need to control ordering
# in the returned dict so that we get nicely ordered yaml files.
spec = {
"name": self.name,
"title": self.title,
"comment": self.comment,
"references": self.references,
"institution": self.institution,
"hierarchical": self.hierarchical,
"last_update": self.last_update.isoformat(),
}
if self.version is not None:
spec["version"] = self.version
spec["total_sum"] = self.total_sum
if self.canonical_top_level_category is not None:
spec["canonical_top_level_category"] = (
self.canonical_top_level_category.codes[0]
)
spec["categories"] = {}
for cat in self.values():
code, cat_spec = cat.to_spec()
spec["categories"][code] = cat_spec
return spec
@functools.cached_property
def _canonical_subgraph(self) -> nx.DiGraph:
return nx.DiGraph(
self._graph.edge_subgraph(
(u, v, 0) for (u, v, s) in self._graph.edges(data="set") if s == 0
)
)
def _show_subtree_children(
self,
children: typing.Iterable[HierarchicalCategory],
format_func: typing.Callable,
prefix: str,
maxdepth: None | int,
) -> str:
children_sorted = natsort.natsorted(children, key=format_func)
r = "".join(
self._show_subtree(
node=child,
prefix=f"{prefix}│",
format_func=format_func,
maxdepth=maxdepth,
)
for child in children_sorted[:-1]
)
# Last child needs to be called slightly differently
r += self._show_subtree(
node=children_sorted[-1],
prefix=f"{prefix} ",
last=True,
format_func=format_func,
maxdepth=maxdepth,
)
return r
@staticmethod
def _render_node(
node: HierarchicalCategory,
last: bool,
prefix: str,
format_func: typing.Callable[[HierarchicalCategory], str],
):
formatted = format_func(node)
if prefix:
if last:
return f"{prefix[:-1]}╰{formatted}\n"
else:
return f"{prefix[:-1]}├{formatted}\n"
else:
return f"{formatted}\n"
def _show_subtree(
self,
*,
node: HierarchicalCategory,
prefix="",
last=False,
format_func: typing.Callable[[HierarchicalCategory], str] = str,
maxdepth: None | int,
) -> str:
"""Recursively-called function to show a subtree starting at the given node."""
r = self._render_node(node, last=last, prefix=prefix, format_func=format_func)
if maxdepth is not None:
maxdepth -= 1
if maxdepth == 0: # maxdepth reached, nothing more to do
return r
child_sets = node.children
if len(child_sets) == 1:
children = child_sets[0]
if children:
r += self._show_subtree_children(
children=children,
format_func=format_func,
maxdepth=maxdepth,
prefix=prefix,
)
elif len(child_sets) > 1:
prefix += "║"
i = 1
for children in child_sets:
if children:
if i == 1:
r += (
f"{prefix[:-1]}╠╤══ ('{format_func(node)}'s children,"
f" option 1)\n"
)
else:
r += (
f"{prefix[:-1]}╠╕ ('{format_func(node)}'s children,"
f" option {i})\n"
)
r += self._show_subtree_children(
children=children,
format_func=format_func,
maxdepth=maxdepth,
prefix=prefix,
)
i += 1
r += f"{prefix[:-1]}╚═══\n"
return r
[docs]
def show_as_tree(
self,
*,
format_func: typing.Callable[[HierarchicalCategory], str] = str,
maxdepth: None | int = None,
root: None | HierarchicalCategory | str = None,
) -> str:
"""Format the hierarchy as a tree.
Starting from the given root, or - if no root is given - the top-level
categories (i.e. categories without parents), the tree of categories that are
transitive children of the root is show, with children connected to their
parents using lines. If a parent category has one set of children, the children
are connected to each other and the parent with a simple line. If a parent
category has multiple sets of children, the sets are connected to parent with
double lines and the children in a set are connected to each other with simple
lines.
Parameters
----------
format_func: callable, optional
Function to call to format categories for display. Each category is
formatted for display using format_func(category), so format_func should
return a string without line breaks, otherwise the tree will look weird.
By default, str() is used, so that the first code and the title of the
category are used.
maxdepth: int, optional
Maximum depth to show in the tree. By default, goes to arbitrary depth.
root: HierarchicalCategory or str, optional
HierarchicalCategory object or code to use as the top-most category.
If not given, the whole tree is shown, starting from all categories without
parents.
Returns
-------
tree_str: str
Representation of the hierarchy as formatted string. print() it for optimal
viewing.
"""
if root is None:
top_level_nodes = (node for node in self.values() if not node.parents)
else:
if not isinstance(root, HierarchicalCategory):
root = self[root]
top_level_nodes = [root]
return "\n".join(
(
self._show_subtree(
node=top_level_node, format_func=format_func, maxdepth=maxdepth
)
)
for top_level_node in top_level_nodes
)
[docs]
def extend(
self,
*,
categories: None | dict[str, dict] = None,
alternative_codes: None | dict[str, str] = None,
children: None | list[tuple] = None,
name: str,
title: None | str = None,
comment: None | str = None,
last_update: None | datetime.date = None,
) -> "HierarchicalCategorization":
"""Extend the categorization with additional categories and relationships,
yielding a new categorization.
Metadata: the ``name``, ``title``, ``comment``, and ``last_update`` are updated
automatically (see below), the ``institution`` and ``references`` are deleted
and the values for ``version``, ``hierarchical``, ``total_sum``, and
``canonical_top_level_category`` are kept.
You can set more accurate metadata (for example, your institution) on the
returned object if needed.
Parameters
----------
categories: dict, optional
Map of new category codes to their specification. The specification is a
dictionary with the keys "title", optionally "comment", and optionally
"alternative_codes".
alternative_codes: dict, optional
Map of new alternative codes. A dictionary with the new alternative code
as key and existing code as value.
children: list, optional
List of ``(parent, (child1, child2, …))`` pairs. The given relationships will
be inserted in the extended categorization.
name : str
The name of your extension. The returned Categorization will have a name
of "{old_name}_{name}", indicating that it is an extension of the underlying
Categorization.
title : str, optional
A string to add to the original title. If not provided, " + {name}" will be
used.
comment : str, optional
A string to add to the original comment. If not provided,
" extended by {name}" will be used.
last_update : datetime.date, optional
The date of the last update to this extension. Today will be used if not
provided.
Returns
-------
Extended categorization : HierarchicalCategorization
"""
spec = self._extend_prepare(
name=name,
categories=categories,
title=title,
comment=comment,
last_update=last_update,
alternative_codes=alternative_codes,
)
if children is not None:
for parent, child_set in children:
if "children" not in spec["categories"][parent]:
spec["categories"][parent]["children"] = []
spec["categories"][parent]["children"].append(child_set)
return HierarchicalCategorization.from_spec(spec)
@property
def df(self) -> "pandas.DataFrame":
"""All category codes as a pandas dataframe."""
titles = []
comments = []
alternative_codes = []
children = []
for cat in self.values():
titles.append(cat.title)
comments.append(cat.comment)
alternative_codes.append(cat.codes[1:])
children.append(
tuple(tuple(sorted(c.codes[0] for c in cs)) for cs in cat.children)
)
return pandas.DataFrame(
index=self.keys(),
data={
"title": titles,
"comment": comments,
"alternative_codes": alternative_codes,
"children": children,
},
)
[docs]
def level(self, cat: str | HierarchicalCategory) -> int:
"""The level of the given category.
The canonical top-level category has level 1 and its children have level 2 etc.
To calculate the level, first only the first ("canonical") set of children is
considered. Only if no path from the canonical top-level category to the
given category can be found all other sets of children are considered to
calculate the level.
"""
if not isinstance(cat, HierarchicalCategory):
return self.level(self[cat])
if not isinstance(self.canonical_top_level_category, HierarchicalCategory):
raise ValueError(
"Can not calculate the level without a canonical_top_level_category."
)
# first use the canonical subgraph for shortest paths
csg = self._canonical_subgraph
try:
sp = nx.shortest_path_length(csg, self.canonical_top_level_category, cat)
except (nx.NetworkXNoPath, nx.NodeNotFound):
try:
sp = nx.shortest_path_length(
self._graph, self.canonical_top_level_category, cat
)
except (nx.NetworkXNoPath, nx.NodeNotFound):
raise ValueError(
f"{cat.codes[0]!r} is not a transitive child of the "
f"canonical top level "
f"{self.canonical_top_level_category.codes[0]!r}."
) from None
return sp + 1
[docs]
def parents(self, cat: str | HierarchicalCategory) -> set[HierarchicalCategory]:
"""The direct parents of the given category."""
if not isinstance(cat, HierarchicalCategory):
return self.parents(self._all_codes_map[cat])
return set(self._graph.predecessors(cat))
[docs]
def ancestors(self, cat: str | HierarchicalCategory) -> set[HierarchicalCategory]:
"""All ancestors of the given category, i.e. the direct parents and their
parents, etc."""
if not isinstance(cat, HierarchicalCategory):
return self.ancestors(self._all_codes_map[cat])
return set(nx.ancestors(self._graph, cat))
[docs]
def children(
self, cat: str | HierarchicalCategory
) -> list[set[HierarchicalCategory]]:
"""The list of sets of direct children of the given category."""
if not isinstance(cat, HierarchicalCategory):
return self.children(self._all_codes_map[cat])
children_dict = {}
for _, child, setno in self._graph.edges(cat, "set"):
if setno not in children_dict:
children_dict[setno] = []
children_dict[setno].append(child)
return [set(children_dict[x]) for x in sorted(children_dict.keys())]
[docs]
def descendants(self, cat: str | HierarchicalCategory) -> set[HierarchicalCategory]:
"""All descendants of the given category, i.e. the direct children and their
children, etc."""
if not isinstance(cat, HierarchicalCategory):
return self.descendants(self._all_codes_map[cat])
return set(nx.descendants(self._graph, cat))
[docs]
def is_leaf(self, cat: str | HierarchicalCategory) -> bool:
"""Is the category a leaf category, i.e. without children?"""
if not isinstance(cat, HierarchicalCategory):
return self.is_leaf(self._all_codes_map[cat])
return cat.is_leaf
[docs]
def leaf_children(
self, cat: str | HierarchicalCategory
) -> list[set[HierarchicalCategory]]:
"""The sets of subcategories which are descendants of the category and do not
have children themselves.
Sets of children are chased separately, so each set of leaf children is
self-sufficient to reconstruct this category (if the categorization allows
reconstructing categories from their children, i.e. if total_sum is set)."""
if not isinstance(cat, HierarchicalCategory):
return self.leaf_children(self._all_codes_map[cat])
return cat.leaf_children
[docs]
def from_pickle(
filepath: str | pathlib.Path | typing.IO[bytes],
) -> Categorization | HierarchicalCategorization:
"""De-serialize Categorization or HierarchicalCategorization from a file written by
to_pickle.
Note that this uses the pickle module, which executes arbitrary code in the
provided file. Only load from pickle files that you trust."""
try:
spec = pickle.load(filepath)
except TypeError:
with open(filepath, "rb") as fd:
spec = pickle.load(fd)
return from_spec(spec)
[docs]
def from_python(
filepath: str | pathlib.Path | typing.IO[bytes],
) -> CategorizationT:
"""Read Categorization or HierarchicalCategorization from a python cache file.
Note that this executes the python cache file. Only load from python cache files
you trust."""
try:
python_code = filepath.read()
filepath.seek(0)
except AttributeError:
python_code = pathlib.Path(filepath).read_text()
variables = {}
exec(python_code, variables)
spec = variables["spec"]
if spec["hierarchical"]:
cls = HierarchicalCategorization
else:
cls = Categorization
return cls.from_spec(spec)
[docs]
def from_spec(spec: dict[str, typing.Any]) -> CategorizationT:
"""Create Categorization or HierarchicalCategorization from a dict specification."""
if spec["hierarchical"]:
return HierarchicalCategorization.from_spec(spec)
else:
return Categorization.from_spec(spec)
[docs]
def from_yaml(
filepath: str | pathlib.Path | typing.TextIO,
) -> CategorizationT:
"""Read Categorization or HierarchicalCategorization from a StrictYaml file."""
try:
yaml = sy.load(filepath.read())
filepath.seek(0)
except AttributeError:
with open(filepath) as fd:
yaml = sy.load(fd.read())
hier = yaml.data["hierarchical"]
if hier in ("yes", "true", "True"):
cls = HierarchicalCategorization
elif hier in ("no", "false", "False"):
cls = Categorization
else:
raise ValueError(
f"'hierarchical' must be 'yes', 'true', 'True', 'no', 'false' or 'False',"
f" not {hier!r}."
)
return cls.from_yaml(filepath)