Commit 1f4263d3 authored by Markus Scheidgen's avatar Markus Scheidgen
Browse files

Refactored annotations and added annotations as fields to MSection.

parent 59a22c73
Pipeline #70189 passed with stages
in 15 minutes and 22 seconds
......@@ -331,9 +331,8 @@ def edit(parsed_query: Dict[str, Any], mongo_update: Dict[str, Any] = None, re_i
if re_index:
def elastic_updates():
for calc in proc.Calc.objects(calc_id__in=calc_ids):
entry = datamodel.EntryMetadata.m_def.m_x('elastic').create_index_entry(
datamodel.EntryMetadata.m_from_dict(calc['metadata']))
entry = entry.to_dict(include_meta=True)
entry_metadata = datamodel.EntryMetadata.m_from_dict(calc['metadata'])
entry = entry_metadata.a_elastic.create_index_entry().to_dict(include_meta=True)
entry['_op_type'] = 'index'
yield entry
......
......@@ -46,7 +46,7 @@ ns = api.namespace(
class CalcMetadata(fields.Raw):
def format(self, value):
entry_metadata = datamodel.EntryMetadata.m_from_dict(value)
return datamodel.EntryMetadata.m_def.m_x('elastic').create_index_entry(entry_metadata).to_dict()
return entry_metadata.a_elastic.create_index_entry().to_dict()
proc_model = api.model('Processing', {
......
......@@ -179,10 +179,8 @@ def index(threads, dry):
with utils.ETA(all_calcs, ' index %10d or %10d calcs, ETA %s') as eta:
for calc in proc.Calc.objects():
eta.add()
entry = None
entry = datamodel.EntryMetadata.m_def.m_x('elastic').create_index_entry(
datamodel.EntryMetadata.m_from_dict(calc.metadata))
entry = entry.to_dict(include_meta=True)
entry_metadata = datamodel.EntryMetadata.m_from_dict(calc.metadata)
entry = entry_metadata.a_elastic.create_index_entry().to_dict(include_meta=True)
entry['_op_type'] = 'index'
yield entry
......
......@@ -16,7 +16,7 @@ from typing import Callable, Any, Dict, cast
import uuid
from .metainfo import Section, Quantity, MSection, Annotation, MEnum, Datetime, Reference
from .metainfo import Section, Quantity, MSection, MEnum, Datetime, Reference, Annotation, SectionAnnotation, DefinitionAnnotation
'''
This module provides metainfo annotation class :class:`Elastic` and
......@@ -25,7 +25,7 @@ metainfo data in elastic search.
'''
class ElasticDocument(Annotation):
class ElasticDocument(SectionAnnotation):
'''
This annotation class can be used to extend metainfo sections. It allows to detail
how section instances (and their sub sections and quantities) should be represented in
......@@ -56,6 +56,9 @@ class ElasticDocument(Annotation):
self.m_def: Section = None
self.fields: Dict[Quantity, str] = {}
def new(self, section):
return dict(elastic=ElasticEntry(section))
def init_annotation(self, definition):
assert isinstance(definition, Section), 'The ElasticDocument annotation is only usable with Sections.'
self.m_def = definition
......@@ -194,7 +197,18 @@ class ElasticDocument(Annotation):
return document
class Elastic(Annotation):
class ElasticEntry(Annotation):
def __init__(self, section: MSection):
self.section = section
def index(self, **kwargs):
return ElasticDocument.index(self.section, **kwargs)
def create_index_entry(self):
return ElasticDocument.create_index_entry(self.section)
class Elastic(DefinitionAnnotation):
'''
This annotation class can be used to extend metainfo quantities. It allows to detail
how this quantity should be represented in an elastic search index.
......
......@@ -656,13 +656,23 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas
MetainfoError('Section has not m_def.')
# get annotations from kwargs
self.m_annotations: Dict[Union[str, type], Any] = {}
rest = {}
self.m_annotations: Dict[str, Any] = {}
other_kwargs = {}
for key, value in kwargs.items():
if key.startswith('a_'):
self.m_annotations[key[2:]] = value
else:
rest[key] = value
other_kwargs[key] = value
# get additional annotations from the section definition
if not is_bootstrapping:
for section_annotation in self.m_def.m_x(SectionAnnotation, as_list=True):
for name, annotation in section_annotation.new(self).items():
self.m_annotations[name] = annotation
# add annotation attributes for names annotations
for annotation_name, annotation in self.m_annotations.items():
setattr(self, 'a_%s' % annotation_name, annotation)
# initialize data
self.m_data = m_data
......@@ -671,9 +681,9 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas
# set remaining kwargs
if is_bootstrapping:
self.m_data.dct.update(**rest) # type: ignore
self.m_data.dct.update(**other_kwargs) # type: ignore
else:
self.m_update(**rest)
self.m_update(**other_kwargs)
@classmethod
def __init_cls__(cls):
......@@ -805,6 +815,12 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas
m_def.__init_metainfo__()
def __getattr__(self, name):
# This will make mypy and pylint ignore 'missing' dynamic attributes and functions
# and wrong types of those.
# Ideally we have a plugin for both that add the corrent type info
return super().__getattr__(name) # pylint: disable=no-member
def __check_np(self, quantity_def: 'Quantity', value: np.ndarray) -> np.ndarray:
# TODO this feels expensive, first check, then possible convert very often?
# if quantity_ref.type != value.dtype:
......@@ -1351,9 +1367,13 @@ class MSection(metaclass=MObjectMeta): # TODO find a way to make this a subclas
return cast(MSectionBound, context)
def m_x(self, key: Union[str, type], default=None, as_list: bool = False):
def m_x(self, *args, **kwargs):
# TODO remove
return self.m_get_annotations(*args, **kwargs)
def m_get_annotations(self, key: Union[str, type], default=None, as_list: bool = False):
'''
Convinience method for get annotations
Convinience method to get annotations
Arguments:
key: Either the optional annoation name or an annotation class. In the first
......@@ -1552,14 +1572,9 @@ class Definition(MSection):
a class context, this method must be called manually on all definitions.
'''
# initialize annotations
for annotation in self.m_annotations.values():
if isinstance(annotation, (tuple, list)):
for single_annotation in annotation:
if isinstance(single_annotation, Annotation):
single_annotation.init_annotation(self)
if isinstance(annotation, Annotation):
annotation.init_annotation(self)
# initialize definition annotations
for annotation in self.m_x(DefinitionAnnotation, as_list=True):
annotation.init_annotation(self)
@classmethod
def all_definitions(cls: Type[MSectionBound]) -> Iterable[MSectionBound]:
......@@ -2066,6 +2081,31 @@ class Category(Definition):
self.definitions: Set[Definition] = set()
class Annotation:
''' Base class for annotations. '''
pass
class DefinitionAnnotation(Annotation):
''' Base class for annotations for definitions. '''
def __init__(self):
self.definition: Definition = None
def init_annotation(self, definition: Definition):
self.definition = definition
class SectionAnnotation(DefinitionAnnotation):
'''
Special annotation class for section definition that allows to auto add annotations
to section instances.
'''
def new(self, section) -> Dict[str, Any]:
return {}
Section.m_def = Section(name='Section')
Section.m_def.m_def = Section.m_def
Section.m_def.section_cls = Section
......@@ -2174,13 +2214,3 @@ class Environment(MSection):
if isinstance(definition, Definition):
definitions = self.all_definitions_by_name.setdefault(definition.name, [])
definitions.append(definition)
class Annotation:
''' Base class for annotations. '''
def __init__(self):
self.definition: Definition = None
def init_annotation(self, definition: Definition):
self.definition = definition
# Copyright 2018 Markus Scheidgen
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an"AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
''' The beginning of a pylint plugin. Unfotunately it is kinda nonsensical without a
partnering mypy plugin. '''
import astroid
from astroid import MANAGER
annotation_names = {
'MSection': {
'a_test': '*'
},
'Section': {
'a_elastic': 'nomad.metainfo.elastic_extension.ElasticDocument'
},
'Quantity': {
'a_elastic': 'nomad.metainfo.elastic_extension.Elastic'
}
}
def register(linter):
# Needed for registering the plugin.
pass
def transform(cls):
for cls_name, annotations in annotation_names.items():
if cls.name == cls_name:
for name, type_spec in annotations.items():
if type_spec == '*':
cls.locals[name] = [astroid.Instance()]
else:
type_path = type_spec.split('.')
type_module = '.'.join(type_path[:-1])
type_name = type_path[-1]
module = MANAGER.ast_from_module_name(type_module)
cls.locals[name] = [cls.instantiate_class() for cls in module.lookup(type_name)[1]]
MANAGER.register_transform(astroid.ClassDef, transform)
......@@ -302,8 +302,7 @@ class Calc(Proc):
entry_metadata.processed = False
self.metadata = entry_metadata.m_to_dict(include_defaults=True)
datamodel.EntryMetadata.m_def.m_x('elastic').index(entry_metadata)
entry_metadata.a_elastic.index()
except Exception as e:
self.get_logger().error('could not index after processing failure', exc_info=e)
......@@ -432,7 +431,7 @@ class Calc(Proc):
# index in search
with utils.timer(logger, 'indexed', step='index'):
datamodel.EntryMetadata.m_def.m_x('elastic').index(entry_metadata)
entry_metadata.a_elastic.index()
# persist the archive
with utils.timer(
......
......@@ -65,7 +65,7 @@ def publish(calcs: Iterable[datamodel.EntryMetadata]) -> None:
''' Update all given calcs with their metadata and set ``publish = True``. '''
def elastic_updates():
for calc in calcs:
entry = calc.m_def.m_x('elastic').create_index_entry(calc)
entry = calc.a_elastic.create_index_entry()
entry.published = True
entry = entry.to_dict(include_meta=True)
source = entry.pop('_source')
......@@ -86,7 +86,7 @@ def index_all(calcs: Iterable[datamodel.EntryMetadata], do_refresh=True) -> None
'''
def elastic_updates():
for calc in calcs:
entry = calc.m_def.m_x('elastic').create_index_entry(calc)
entry = calc.a_elastic.create_index_entry()
entry = entry.to_dict(include_meta=True)
entry['_op_type'] = 'index'
yield entry
......
......@@ -717,7 +717,7 @@ class TestRepo():
entry_metadata.m_update(
calc_id='1', uploader=test_user.user_id, published=True, with_embargo=False)
EntryMetadata.m_def.m_x('elastic').index(entry_metadata, refresh=True)
entry_metadata.a_elastic.index(refresh=True)
entry_metadata.m_update(
calc_id='2', uploader=other_test_user.user_id, published=True,
......@@ -726,17 +726,17 @@ class TestRepo():
entry_metadata.m_update(
atoms=['Fe'], comment='this is a specific word', formula='AAA')
entry_metadata.dft.basis_set = 'zzz'
EntryMetadata.m_def.m_x('elastic').index(entry_metadata, refresh=True)
entry_metadata.a_elastic.index(refresh=True)
entry_metadata.m_update(
calc_id='3', uploader=other_test_user.user_id, published=False,
with_embargo=False, pid=3, external_id='external_3')
EntryMetadata.m_def.m_x('elastic').index(entry_metadata, refresh=True)
entry_metadata.a_elastic.index(refresh=True)
entry_metadata.m_update(
calc_id='4', uploader=other_test_user.user_id, published=True,
with_embargo=True, pid=4, external_id='external_4')
EntryMetadata.m_def.m_x('elastic').index(entry_metadata, refresh=True)
entry_metadata.a_elastic.index(refresh=True)
yield
......@@ -1798,7 +1798,7 @@ class TestDataset:
Calc(
calc_id='1', upload_id='1', create_time=datetime.datetime.now(),
metadata=entry_metadata.m_to_dict()).save()
EntryMetadata.m_def.m_x('elastic').index(entry_metadata, refresh=True)
entry_metadata.a_elastic.index(refresh=True)
def test_delete_dataset(self, api, test_user_auth, example_dataset_with_entry):
rv = api.delete('/datasets/ds1', headers=test_user_auth)
......
......@@ -698,6 +698,6 @@ def create_test_structure(
proc_calc = processing.Calc.from_entry_metadata(calc)
proc_calc.save()
calc.m_def.m_x('elastic').index(calc)
calc.a_elastic.index()
assert processing.Calc.objects(calc_id__in=[calc.calc_id]).count() == 1
......@@ -21,7 +21,8 @@ from nomadcore.local_meta_info import InfoKindEl, InfoKindEnv
from nomad.metainfo.metainfo import (
MSection, MCategory, Section, Quantity, SubSection, Definition, Package, DeriveError,
MetainfoError, Environment, MResource, Datetime, units, Annotation)
MetainfoError, Environment, MResource, Datetime, units, Annotation, SectionAnnotation,
DefinitionAnnotation)
from nomad.metainfo.example import Run, VaspRun, System, SystemHash, Parsing, m_package as example_package
from nomad.metainfo.legacy import LegacyMetainfoEnvironment
from nomad.parsing.metainfo import MetainfoBackend
......@@ -237,19 +238,24 @@ class TestM2:
assert System.n_atoms.virtual
def test_annotations(self):
class TestSectionAnnotation(Annotation):
class TestSectionAnnotation(SectionAnnotation):
def init_annotation(self, definition):
super().init_annotation(definition)
section_cls = definition.section_cls
assert definition.name == 'TestSection'
assert 'test_quantity' in definition.all_quantities
assert definition.all_quantities['test_quantity'].m_x('test').initialized
assert definition.all_quantities['test_quantity'].m_x('test', as_list=True)[0].initialized
assert definition.all_quantities['test_quantity'].m_x(Annotation).initialized
assert all(a.initialized for a in definition.all_quantities['list_test_quantity'].m_x('test'))
assert all(a.initialized for a in definition.all_quantities['list_test_quantity'].m_x(Annotation))
assert section_cls.test_quantity.m_get_annotations('test').initialized
assert section_cls.test_quantity.a_test.initialized
assert section_cls.test_quantity.m_get_annotations('test', as_list=True)[0].initialized
assert section_cls.test_quantity.m_get_annotations(Annotation).initialized
assert all(a.initialized for a in section_cls.list_test_quantity.a_test)
assert all(a.initialized for a in section_cls.list_test_quantity.m_get_annotations(Annotation))
self.initialized = True
class TestQuantityAnnotation(Annotation):
def new(self, section):
return dict(test='test annotation')
class TestQuantityAnnotation(DefinitionAnnotation):
def init_annotation(self, definition):
super().init_annotation(definition)
assert definition.name in ['test_quantity', 'list_test_quantity']
......@@ -264,8 +270,10 @@ class TestM2:
type=str,
a_test=[TestQuantityAnnotation(), TestQuantityAnnotation()])
assert TestSection.m_def.m_x('test').initialized
assert TestSection.m_def.m_x(TestSectionAnnotation).initialized
assert TestSection.m_def.a_test.initialized
assert TestSection.m_def.m_get_annotations(TestSectionAnnotation).initialized
assert TestSection().a_test == 'test annotation'
class TestM1:
......
......@@ -150,7 +150,7 @@ if __name__ == '__main__':
with upload_files.archive_log_file(calc.calc_id, 'wt') as f:
f.write('this is a generated test file')
search_entry = calc.m_def.m_x('elastic').create_index_entry(calc)
search_entry = calc.a_elastic.create_index_entry()
search_entry.n_total_energies = random.choice(low_numbers_for_total_energies)
search_entry.n_geometries = low_numbers_for_geometries
for _ in range(0, random.choice(search_entry.n_geometries)):
......
......@@ -233,8 +233,7 @@ def refresh_index():
def create_entry(entry_metadata: datamodel.EntryMetadata):
entry = datamodel.EntryMetadata.m_def.m_x('elastic').index(entry_metadata)
entry.save()
entry = entry_metadata.a_elastic.index()
assert_entry(entry_metadata.calc_id)
return entry
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment