followthemoney 1.3.7__py3-none-any.whl → 3.8.0__py3-none-any.whl

followthemoney/mapping/property.py

@@ -1,46 +1,76 @@
 import re
 from copy import deepcopy
+from warnings import warn
 from normality import stringify
-from banal import unique_list, ensure_list
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
+from banal import keys_values, as_bool
 
+from followthemoney.helpers import inline_names
 from followthemoney.exc import InvalidMapping
+from followthemoney.proxy import EntityProxy
+from followthemoney.util import sanitize_text
+from followthemoney.property import Property
+from followthemoney.mapping.source import Record
+
+if TYPE_CHECKING:
+    from followthemoney.mapping.query import QueryMapping
 
 
 class PropertyMapping(object):
     """Map values from a given record (e.g. a CSV row or SQL result) to the
     schema form."""
-    FORMAT_PATTERN = re.compile('{{([^(}})]*)}}')
 
-    def __init__(self, query, data, schema):
+    __slots__ = (
+        "query",
+        "prop",
+        "refs",
+        "join",
+        "split",
+        "entity",
+        "format",
+        "fuzzy",
+        "required",
+        "literals",
+        "template",
+        "replacements",
+    )
+
+    FORMAT_PATTERN = re.compile("{{([^(}})]*)}}")
+
+    def __init__(
+        self, query: "QueryMapping", data: Dict[str, Any], prop: Property
+    ) -> None:
         self.query = query
         data = deepcopy(data)
-        self.data = data
-        self.schema = schema
-        self.name = schema.name
-        self.type = schema.type
-
-        self.refs = ensure_list(data.pop('column', []))
-        self.refs.extend(ensure_list(data.pop('columns', [])))
-
-        self.literals = ensure_list(data.pop('literal', []))
-        self.literals.extend(ensure_list(data.pop('literals', [])))
-
-        self.join = data.pop('join', None)
-        self.split = data.pop('split', None)
-        self.entity = data.pop('entity', None)
-        self.required = data.pop('required', False)
-
-        self.template = stringify(data.pop('template', None))
-        self.replacements = {}
+        self.prop = prop
+
+        self.refs = cast(List[str], keys_values(data, "column", "columns"))
+        self.join = cast(Optional[str], data.pop("join", None))
+        self.split = cast(Optional[str], data.pop("split", None))
+        self.entity = stringify(data.pop("entity", None))
+        self.format = stringify(data.pop("format", None))
+        self.fuzzy = as_bool(data.pop("fuzzy", False))
+        self.required = as_bool(data.pop("required", False))
+        self.literals = cast(List[str], keys_values(data, "literal", "literals"))
+
+        self.template = sanitize_text(data.pop("template", None))
+        self.replacements: Dict[str, str] = {}
         if self.template is not None:
             # this is hacky, trying to generate refs from template
             for ref in self.FORMAT_PATTERN.findall(self.template):
                 self.refs.append(ref)
-                self.replacements['{{%s}}' % ref] = ref
+                self.replacements["{{%s}}" % ref] = ref
 
-    def bind(self):
-        if self.schema.stub:
-            raise InvalidMapping("Property for [%s] is a stub" % self.name)
+    def bind(self) -> None:
+        if self.prop.stub:
+            raise InvalidMapping("Property for [%r] is a stub" % self.prop)
+
+        if self.prop.deprecated:
+            warn(
+                "Mapping uses a deprecated property: %r" % self.prop,
+                DeprecationWarning,
+                stacklevel=2,
+            )
 
         if self.entity is None:
             return
@@ -52,44 +82,47 @@ class PropertyMapping(object):
         for entity in self.query.entities:
             if entity.name != self.entity:
                 continue
-            if not entity.schema.is_a(self.schema.range):
-                raise InvalidMapping("The entity [%s] must be a %s (not %s)" %
-                                     (self.name, self.schema.range, entity.schema.name))  # noqa
+            if not self.prop.range or not entity.schema.is_a(self.prop.range):
+                raise InvalidMapping(
+                    "The entity [%r] must be a %s (not %s)"
+                    % (self.prop, self.prop.range, entity.schema.name)
+                )  # noqa
             return
 
-        raise InvalidMapping("No entity [%s] for property [%s]"
-                             % (self.entity, self.name))
+        raise InvalidMapping(
+            "No entity [%s] for property [%r]" % (self.entity, self.prop)
+        )
 
-    def record_values(self, record):
+    def record_values(self, record: Record) -> List[str]:
         if self.template is not None:
             # replace mentions of any refs with the values present in the
             # current record
             value = self.template
             for repl, ref in self.replacements.items():
-                ref_value = record.get(ref) or ''
+                ref_value = record.get(ref) or ""
                 value = value.replace(repl, ref_value)
             return [value.strip()]
 
         values = list(self.literals)
-        values.extend([record.get(r) for r in self.refs])
+        for ref in self.refs:
+            rec_value = record.get(ref)
+            if rec_value is not None:
+                values.append(rec_value)
         return values
 
-    def map(self, record, entities, **kwargs):
-        kwargs.update(self.data)
-
+    def map(
+        self, proxy: EntityProxy, record: Record, entities: Dict[str, EntityProxy]
+    ) -> List[str]:
         if self.entity is not None:
             entity = entities.get(self.entity)
             if entity is not None:
-                return [entity.id]
+                proxy.unsafe_add(self.prop, entity.id, cleaned=True)
+                inline_names(proxy, entity)
             return []
 
         # clean the values returned by the query, or by using literals, or
         # formats.
-        values = []
-        for value in self.record_values(record):
-            value = self.type.clean(value, **kwargs)
-            if value is not None:
-                values.append(value)
+        values: List[str] = self.record_values(record)
 
         if self.join is not None:
             values = [self.join.join(values)]
@@ -97,7 +130,20 @@ class PropertyMapping(object):
         if self.split is not None:
             splote = []
             for value in values:
-                splote = splote + value.split(self.split)
+                splote.extend(value.split(self.split))
             values = splote
 
-        return unique_list(values)
+        discarded_values: List[str] = []
+
+        for value in values:
+            added_value = proxy.unsafe_add(
+                prop=self.prop,
+                value=value,
+                fuzzy=self.fuzzy,
+                format=self.format,
+            )
+
+            if value is not None and added_value is None:
+                discarded_values.append(value)
+
+        return discarded_values

followthemoney/mapping/query.py

@@ -1,20 +1,27 @@
+from followthemoney.mapping.source import Record, Source
+from typing import TYPE_CHECKING, Any, List, Optional, Set, Dict
+
+from followthemoney.proxy import EntityProxy
 from followthemoney.mapping.entity import EntityMapping
 from followthemoney.mapping.sql import SQLSource
 from followthemoney.mapping.csv import CSVSource
 from followthemoney.exc import InvalidMapping
 
+if TYPE_CHECKING:
+    from followthemoney.model import Model
 
-class QueryMapping(object):
 
-    def __init__(self, model, data, key_prefix=None):
-        self.model = model
-        self.data = data
+class QueryMapping:
+    __slots__ = ("model", "data", "refs", "entities", "source")
 
-        self.refs = set()
-        self.entities = []
-        for name, data in data.get('entities', {}).items():
-            entity = EntityMapping(model, self, name, data,
-                                   key_prefix=key_prefix)
+    def __init__(
+        self, model: "Model", data: Dict[str, Any], key_prefix: Optional[str] = None
+    ) -> None:
+        self.model = model
+        self.refs: Set[str] = set()
+        self.entities: List[EntityMapping] = []
+        for name, edata in data.get("entities", {}).items():
+            entity = EntityMapping(model, self, name, edata, key_prefix=key_prefix)
 
             self.entities.append(entity)
             self.refs.update(entity.refs)
@@ -32,7 +39,7 @@ class QueryMapping(object):
         # in dependent entities.
         entities = self.entities
         self.entities = []
-        resolved = set()
+        resolved: Set[str] = set()
         while len(entities) > 0:
             before = len(entities)
             for entity in entities:
@@ -44,16 +51,17 @@ class QueryMapping(object):
             if before == len(entities):
                 raise InvalidMapping("Circular entity dependency detected.")
 
-    @property
-    def source(self):
-        if 'database' in self.data:
-            return SQLSource(self, self.data)
-        elif 'csv_url' in self.data or 'csv_urls' in self.data:
-            return CSVSource(self, self.data)
-        raise InvalidMapping("Cannot determine mapping type")
+        self.source = self._get_source(data)
+
+    def _get_source(self, data: Dict[str, Any]) -> Source:
+        if "database" in data:
+            return SQLSource(self, data)
+        if "csv_url" in data or "csv_urls" in data:
+            return CSVSource(self, data)
+        raise InvalidMapping("Cannot determine mapping type: %r" % data)
 
-    def map(self, record):
-        data = {}
+    def map(self, record: Record) -> Dict[str, EntityProxy]:
+        data: Dict[str, EntityProxy] = {}
         for entity in self.entities:
             mapped = entity.map(record, data)
             if mapped is not None:

followthemoney/mapping/source.py

@@ -1,11 +1,21 @@
+from typing import TYPE_CHECKING, Any, Dict, Generator, Optional, Set, cast
 
+if TYPE_CHECKING:
+    from followthemoney.mapping.query import QueryMapping
+
+Filter = Set[Optional[str]]
+Record = Dict[str, str]
 
-class Source(object):
 
-    def __init__(self, query, data):
+class Source(object):
+    def __init__(self, query: "QueryMapping", data: Dict[str, Any]) -> None:
         self.query = query
-        self.filters = query.data.get('filters', {}).items()
-        self.filters_not = query.data.get('filters_not', {}).items()
+        self.filters = cast(Dict[str, Any], data.get("filters", {})).items()
+        self.filters_not = cast(Dict[str, Any], data.get("filters_not", {})).items()
+
+    @property
+    def records(self) -> Generator[Record, None, None]:
+        raise NotImplementedError
 
-    def __len__(self):
+    def __len__(self) -> int:
         return 0

followthemoney/mapping/sql.py

@@ -1,18 +1,24 @@
 import os
-import six
 import logging
 from uuid import uuid4
-from banal import ensure_list
-from normality import stringify
-from sqlalchemy import create_engine, MetaData
-from sqlalchemy import select, func
-# from sqlalchemy import text as sql_text
+from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Union, cast
+from banal import ensure_list, is_listish, keys_values
+from sqlalchemy import MetaData, func
+from sqlalchemy.future import select
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.sql.elements import Label
 from sqlalchemy.pool import NullPool
 from sqlalchemy.schema import Table
+from sqlalchemy.sql.expression import Select
 
-from followthemoney.mapping.source import Source
+from followthemoney.mapping.source import Record, Source
+from followthemoney.util import sanitize_text
 from followthemoney.exc import InvalidMapping
 
+if TYPE_CHECKING:
+    from followthemoney.mapping.query import QueryMapping
+
+
 log = logging.getLogger(__name__)
 DATA_PAGE = 1000
 
@@ -20,20 +26,22 @@ DATA_PAGE = 1000
 class QueryTable(object):
     """A table to be joined in."""
 
-    def __init__(self, query, data):
-        self.query = query
-        if isinstance(data, six.string_types):
-            data = {'table': data}
-        self.data = data
-        self.table_ref = data.get('table')
-        self.alias_ref = data.get('alias', self.table_ref)
-        self.table = Table(self.table_ref, self.query.meta, autoload=True)
-        self.alias = self.table.alias(self.alias_ref)
-
-        self.refs = {}
+    def __init__(
+        self, meta: MetaData, engine: Engine, data: Union[str, Dict[str, str]]
+    ) -> None:
+        if isinstance(data, str):
+            data = {"table": data}
+        table_ref = data.get("table")
+        if table_ref is None:
+            raise InvalidMapping("Query has no table!")
+        alias_ref = data.get("alias", table_ref)
+        self.table = Table(table_ref, meta, autoload_with=engine)
+        self.alias = self.table.alias(alias_ref)
+
+        self.refs: Dict[str, Label[Any]] = {}
         for column in self.alias.columns:
-            name = '%s.%s' % (self.alias_ref, column.name)
-            labeled_column = column.label('col_%s' % uuid4().hex[:10])
+            name = "%s.%s" % (alias_ref, column.name)
+            labeled_column = column.label("col_%s" % uuid4().hex[:10])
             self.refs[name] = labeled_column
             self.refs[column.name] = labeled_column
 
@@ -41,70 +49,76 @@ class QueryTable(object):
 class SQLSource(Source):
     """Query mapper for loading data from a SQL query."""
 
-    def __init__(self, query, data):
+    def __init__(self, query: "QueryMapping", data: Dict[str, Any]) -> None:
         super(SQLSource, self).__init__(query, data)
-        self.database_uri = os.path.expandvars(data.get('database'))
-        kwargs = {}
-        if self.database_uri.lower().startswith('postgres'):
-            kwargs['server_side_cursors'] = True
-        self.engine = create_engine(self.database_uri,
-                                    poolclass=NullPool,
-                                    **kwargs)
+        database = data.get("database")
+        if database is None:
+            raise InvalidMapping("No database in SQL mapping!")
+        self.database_uri = cast(str, os.path.expandvars(database))
+        self.engine = create_engine(self.database_uri, poolclass=NullPool)
         self.meta = MetaData()
-        self.meta.bind = self.engine
 
-        tables = ensure_list(data.get('table'))
-        tables.extend(ensure_list(data.get('tables')))
-        self.tables = [QueryTable(self, f) for f in tables]
-        self.joins = ensure_list(data.get('joins'))
+        tables = keys_values(data, "table", "tables")
+        self.tables = [QueryTable(self.meta, self.engine, f) for f in tables]
+        self.joins = cast(List[Dict[str, str]], ensure_list(data.get("joins")))
 
-    def get_column(self, ref):
+    def get_column(self, ref: Optional[str]) -> Label[Any]:
         for table in self.tables:
             if ref in table.refs:
-                return table.refs.get(ref)
+                return table.refs[ref]
         raise InvalidMapping("Missing reference: %s" % ref)
 
-    def apply_filters(self, q):
+    def apply_filters(self, q: Select) -> Select:
         for col, val in self.filters:
-            q = q.where(self.get_column(col) == val)
+            if is_listish(val):
+                q = q.where(self.get_column(col).in_(val))
+            else:
+                q = q.where(self.get_column(col) == val)
         for col, val in self.filters_not:
-            q = q.where(self.get_column(col) != val)
+            if is_listish(val):
+                q = q.where(self.get_column(col).notin_(val))
+            else:
+                q = q.where(self.get_column(col) != val)
         # not sure this is a great idea:
         # if self.data.get('where'):
         #    q = q.where(sql_text(self.data.get('where')))
         for join in self.joins:
-            left = self.get_column(join.get('left'))
-            right = self.get_column(join.get('right'))
+            left = self.get_column(join.get("left"))
+            right = self.get_column(join.get("right"))
             q = q.where(left == right)
         return q
 
-    def compose_query(self):
-        from_clause = [t.alias for t in self.tables]
+    def compose_query(self) -> Select:
         columns = [self.get_column(r) for r in self.query.refs]
-        q = select(columns=columns, from_obj=from_clause, use_labels=True)
+        q = select(*columns)
+        q = q.select_from(*[t.alias for t in self.tables])
         return self.apply_filters(q)
 
     @property
-    def records(self):
+    def records(self) -> Generator[Record, None, None]:
         """Compose the actual query and return an iterator of ``Record``."""
         mapping = [(r, self.get_column(r).name) for r in self.query.refs]
         q = self.compose_query()
         log.info("Query: %s", q)
-        rp = self.engine.execute(q)
-        while True:
-            rows = rp.fetchmany(size=DATA_PAGE)
-            if not len(rows):
-                break
-            for row in rows:
-                data = {}
-                for ref, name in mapping:
-                    data[ref] = stringify(row[name])
-                yield data
-
-    def __len__(self):
-        from_clause = [t.alias for t in self.tables]
-        columns = [func.count('*')]
-        q = select(columns=columns, from_obj=from_clause, use_labels=True)
+        with self.engine.connect() as conn:
+            rp = conn.execution_options(stream_results=True).execute(q)
+            while True:
+                rows = rp.fetchmany(size=DATA_PAGE)
+                if not len(rows):
+                    break
+                for row in rows:
+                    row_map = row._mapping
+                    data: Record = {}
+                    for ref, name in mapping:
+                        value = sanitize_text(row_map[name])
+                        if value is not None:
+                            data[ref] = value
+                    yield data
+
+    def __len__(self) -> int:
+        q = select(func.count("*"))
+        q = q.select_from(*[t.alias for t in self.tables])
         q = self.apply_filters(q)
-        rp = self.engine.execute(q)
-        return rp.scalar()
+        with self.engine.connect() as conn:
+            rp = conn.execute(q)
+            return int(rp.scalar() or 0)

followthemoney/messages.py

@@ -1,19 +1,25 @@
-import six
 import yaml
+from typing import Any, Dict, Generator, List, TextIO, Tuple
 
+Message = Tuple[Any, Any, List[str], List[str]]
 
-def extract_object(data, path):
+
+def extract_object(
+    data: Dict[str, Any], path: List[str]
+) -> Generator[Message, None, None]:
     for key, value in data.items():
         subpath = path + [key]
-        if isinstance(value, six.string_types):
-            if key in ['label', 'reverse', 'description', 'plural']:
-                comment = '.'.join(subpath)
+        if isinstance(value, str):
+            if key in ["label", "reverse", "description", "plural"]:
+                comment = ".".join(subpath)
                 yield (None, None, [value], [comment])
         if isinstance(value, dict):
             for res in extract_object(value, subpath):
                 yield res
 
 
-def extract_yaml(fileobj, keywords, comment_tags, options):
-    data = yaml.load(fileobj)
+def extract_yaml(
+    fileobj: TextIO, keywords: Any, comment_tags: Any, options: Any
+) -> Generator[Message, None, None]:
+    data = yaml.safe_load(fileobj)
     return extract_object(data, [])