Repository: EmilStenstrom/conllu

Issue:
A PEP 561 compliant package
First off, thanks for this package. The API on the README looks great. Can't wait to dig in.

I've forked this repo and will add type annotations to every part that I will make use(most likely all the `.py` files in the root dir). 

If you'd like I can submit a PR to merge those type annotations back in(and this repo can be PEP 561 compliant).

Discussion:
davidatbu: I am having trouble parsing the CI errors, would be happy to correct if you point me in the right direction @EmilStenstrom 
davidatbu: Ok, I think I've figued out why the CI is failing. The only thing that is failing is code coverage. And it's because of two changes that I made that testing doesn't cover:

```py
import typing as T

if T.TYPE_CHECKING:
   ...
```

This one is necessary to avoid cyclical imports _and_ annotate the argument type properly.

The other one is:

```py
class TokenTree:
    def __eq__(self, other: T.Any) -> bool:
        if isinstance(other, TokenTree):
            return self.token == other.token and self.children == other.children \
                and self.metadata == other.metadata
        return False
```

Before, the code assumed that `other` would also be of type `TokenTree`, but actually, `__eq__` needs to support comparing to other types of objects as well(mypy complains if one changes the annotation of `other` to `: 'TokenTree'`).

The second error can be addressed by adding a unit test. The first error, I'm not sure how to deal with, except for configuring the test coverage tool to not care about `if T.TYPE_CHECKING`, which I've no idea is possible or not.

EmilStenstrom: The Typing error could be fixed by simply ignoring that line. I think it could be done by adding “  # pragma: no cover” after the typing if, or maybe adding if T.Typing here: https://github.com/EmilStenstrom/conllu/search?q=pragma&type=
EmilStenstrom: @davidatbu I've released version 4.0 of conllu, that drops support for anything < Python 3.6. If you rebase from latest master you should approach a shippable version of this PR.
EmilStenstrom: Let's not add black right now, we can look at that after this fix has landed.
EmilStenstrom: So next steps:

1. Rebase this branch on latest master
2. Fix any new typing errors reported by mypy
3. Add mypy checks to tox.ini and make sure they pass on CI
4. Refactor TokenList.extend to only support extending with TokenList or "Lists of Dicts" according to my comment above. 

Does this make sense?

Patch:
diff --git a/conllu/__init__.py b/conllu/__init__.py
index 8fad54b..3c48a8b 100644
--- a/conllu/__init__.py
+++ b/conllu/__init__.py
@@ -1,12 +1,18 @@
 from __future__ import unicode_literals
 
+import typing as T
 from io import StringIO
 
-from conllu.models import TokenList
-from conllu.parser import parse_conllu_plus_fields, parse_sentences, parse_token_and_metadata
+from conllu.models import TokenList, TokenTree
+from conllu.parser import (
+    _FieldParserType, _MetadataParserType, parse_conllu_plus_fields, parse_sentences, parse_token_and_metadata,
+)
 
 
-def parse(data, fields=None, field_parsers=None, metadata_parsers=None):
+def parse(data: str, fields: T.Optional[T.Sequence[str]] = None,
+          field_parsers: T.Dict[str, _FieldParserType] = None,
+          metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None
+          ) -> T.List[TokenList]:
     return list(parse_incr(
         StringIO(data),
         fields=fields,
@@ -14,7 +20,10 @@ def parse(data, fields=None, field_parsers=None, metadata_parsers=None):
         metadata_parsers=metadata_parsers
     ))
 
-def parse_incr(in_file, fields=None, field_parsers=None, metadata_parsers=None):
+def parse_incr(in_file: T.TextIO, fields: T.Optional[T.Sequence[str]] = None,
+               field_parsers: T.Dict[str, _FieldParserType] = None,
+               metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None
+               ) -> T.Iterator[TokenList]:
     if not hasattr(in_file, 'read'):
         raise FileNotFoundError("Invalid file, 'parse_incr' needs an opened file as input")
 
@@ -29,9 +38,9 @@ def parse_incr(in_file, fields=None, field_parsers=None, metadata_parsers=None):
             metadata_parsers=metadata_parsers
         ))
 
-def parse_tree(data):
+def parse_tree(data: str) -> T.List[TokenTree]:
     return list(parse_tree_incr(StringIO(data)))
 
-def parse_tree_incr(in_file):
+def parse_tree_incr(in_file: T.TextIO) -> T.Iterator[TokenTree]:
     for tokenlist in parse_incr(in_file):
         yield tokenlist.to_tree()
diff --git a/conllu/models.py b/conllu/models.py
index 339bb61..7869611 100644
--- a/conllu/models.py
+++ b/conllu/models.py
@@ -1,6 +1,8 @@
 from __future__ import print_function, unicode_literals
 
+import typing as T
 from collections import OrderedDict, defaultdict
+from collections.abc import Mapping
 
 from conllu.exceptions import ParseException
 from conllu.serializer import serialize
@@ -18,70 +20,65 @@ class Token(OrderedDict):
         "xpos": "xpostag",
     }
 
-    def get(self, key, *args, **kwargs):
+    def get(self, key: str, default: T.Optional[T.Any] = None) -> T.Any:
         if key not in self and key in self.MAPPING:
             key = self.MAPPING[key]
 
-        return super(Token, self).get(key, *args, **kwargs)
+        return super(Token, self).get(key, default)
 
-    def __missing__(self, key):
+    def __missing__(self, key: str) -> T.Any:
         if key in self.MAPPING:
             return self[self.MAPPING[key]]
 
         raise KeyError("'" + key + "'")
 
-class TokenList(list):
-    metadata = None
+class TokenList(T.List[Token]):
 
-    def __init__(self, tokens, metadata=None):
+    metadata: Metadata = Metadata()
+
+    def __init__(self, tokens: T.Iterable[Token], metadata: Metadata = None):
         super(TokenList, self).__init__(tokens)
         if not isinstance(tokens, list):
             raise ParseException("Can't create TokenList, tokens is not a list.")
 
-        self.metadata = metadata
+        self.metadata = metadata or Metadata()
 
-    def __repr__(self):
-        return 'TokenList<' + ', '.join(token['form'] for token in self) + '>'
+    def __repr__(self) -> str:
+        return 'TokenList<' + ', '.join(token['form'] for token in self if 'form' in token) + '>'
 
-    def __eq__(self, other):
+    def __eq__(self, other: T.Any) -> bool:
         return super(TokenList, self).__eq__(other) \
             and (not hasattr(other, 'metadata') or self.metadata == other.metadata)
 
-    def __ne__(self, other):
+    def __ne__(self, other: T.Any) -> bool:
         return not self == other
 
-    def clear(self):
+    def clear(self) -> None:
         self[:] = []  # Supported in Python 2 and 3, unlike clear()
-        self.metadata = None
+        self.metadata = Metadata()
 
-    def copy(self):
+    def copy(self) -> 'TokenList':
         tokens_copy = self[:]  # Supported in Python 2 and 3, unlike copy()
         return TokenList(tokens_copy, self.metadata)
 
-    def extend(self, iterable):
+    def extend(self, iterable: T.Union['TokenList', T.Iterable[Token]]) -> None:
         super(TokenList, self).extend(iterable)
-        if hasattr(iterable, 'metadata'):
-            if hasattr(self.metadata, '__add__') and hasattr(iterable.metadata, '__add__'):
-                self.metadata += iterable.metadata
-            elif type(self.metadata) is dict and type(iterable.metadata) is dict:
-                # noinspection PyUnresolvedReferences
-                self.metadata.update(iterable.metadata)
-            else:
-                self.metadata = [self.metadata, iterable.metadata]
+        if isinstance(iterable, TokenList):
+            self.metadata.update(iterable.metadata)
 
     @property
-    def tokens(self):
+    def tokens(self) -> 'TokenList':
         return self
 
     @tokens.setter
-    def tokens(self, value):
+    def tokens(self, value: T.Iterable[Token]) -> None:
         self[:] = value  # Supported in Python 2 and 3, unlike clear()
 
-    def serialize(self):
+    def serialize(self) -> str:
         return serialize(self)
 
     @staticmethod
-    def head_to_token(sentence):
+    def head_to_token(sentence: 'TokenList') -> T.Dict[int, T.List[Token]]:
         if not sentence:
             raise ParseException("Can't parse tree, need a tokenlist as input.")
 
@@ -109,8 +106,8 @@ class TokenList(list):
 
         return head_indexed
 
-    def to_tree(self):
-        def _create_tree(head_to_token_mapping, id_=0):
+    def to_tree(self) -> 'TokenTree':
+        def _create_tree(head_to_token_mapping: T.Dict[int, T.List[Token]], id_: int = 0) -> T.List['TokenTree']:
             return [
                 TokenTree(child, _create_tree(head_to_token_mapping, child["id"]))
                 for child in head_to_token_mapping[id_]
@@ -120,8 +117,8 @@ class TokenList(list):
         root.set_metadata(self.metadata)
         return root
 
-    def filter(self, **kwargs):
-        tokens = self.tokens.copy()
+    def filter(self, **kwargs: T.Any) -> 'TokenList':
+        tokens: T.Iterable[Token] = self.tokens.copy()
 
         for query, value in kwargs.items():
             filtered_tokens = []
@@ -133,48 +130,56 @@ class TokenList(list):
 
         return TokenList(tokens)
 
-def traverse_dict(obj, query):
+
+_T = T.TypeVar("_T")
+def traverse_dict(obj: T.Mapping[str, _T], query: str) -> T.Optional[_T]:
     """
         Get elements inside a nested dict, based on a dict query. The query is defined by a
         string separated by '__'. traverse_dict(foo, 'a__b__c') is roughly equivalent to foo[a][b][c] but
         will short circuit to return None if something on the query is None.
     """
-    query = query.split('__')
-    for name in query:
-        obj = obj.get(name, None)
-        if obj is None:
+    query_split = query.split('__')
+
+    cur_obj: T.Optional[T.Union[_T, T.Mapping[str, _T]]] = obj
+    for name in query_split:
+        assert isinstance(cur_obj, Mapping)  # help mypy
+        cur_obj = cur_obj.get(name, None)
+        if cur_obj is None:
             return None
-    return obj
+    assert not isinstance(cur_obj, Mapping)  # help mypy
+    return cur_obj
 
 
 class TokenTree(object):
-    token = None
-    children = None
-    metadata = None
+    token: Token
+    children: T.List['TokenTree']
+    metadata: T.Optional[Metadata]
 
-    def __init__(self, token, children, metadata=None):
+    def __init__(self, token: Token, children: T.List['TokenTree'], metadata: T.Optional[Metadata] = None):
         self.token = token
         self.children = children
         self.metadata = metadata
 
-    def set_metadata(self, metadata):
+    def set_metadata(self, metadata: T.Optional[Metadata]) -> None:
         self.metadata = metadata
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return 'TokenTree<' + \
-            'token={id=' + str(self.token['id']) + ', form=' + self.token['form'] + '}, ' + \
+            'token={id=' + str(self.token['id']) + ', form=' + str(self.token['form']) + '}, ' + \
             'children=' + ('[...]' if self.children else 'None') + \
             '>'
 
-    def __eq__(self, other):
-        return self.token == other.token and self.children == other.children \
-            and self.metadata == other.metadata
+    def __eq__(self, other: T.Any) -> bool:
+        if isinstance(other, TokenTree):
+            return self.token == other.token and self.children == other.children \
+                and self.metadata == other.metadata
+        return False
 
-    def serialize(self):
+    def serialize(self) -> str:
         if not self.token or "id" not in self.token:
             raise ParseException("Could not serialize tree, missing 'id' field.")
 
-        def flatten_tree(root_token, token_list=[]):
+        def flatten_tree(root_token: TokenTree, token_list: T.List[Token] = []) -> T.List[Token]:
             token_list.append(root_token.token)
 
             for child_token in root_token.children:
@@ -188,7 +193,8 @@ class TokenTree(object):
 
         return serialize(tokenlist)
 
-    def print_tree(self, depth=0, indent=4, exclude_fields=DEFAULT_EXCLUDE_FIELDS):
+    def print_tree(self, depth: int = 0, indent: int = 4,
+                   exclude_fields: T.Sequence[str] = DEFAULT_EXCLUDE_FIELDS) -> None:
         if not self.token:
             raise ParseException("Can't print, token is None.")
 
diff --git a/conllu/parser.py b/conllu/parser.py
index b8050e4..1f7fa7e 100644
--- a/conllu/parser.py
+++ b/conllu/parser.py
@@ -1,13 +1,18 @@
 from __future__ import unicode_literals
 
 import re
+import typing as T
 from collections import OrderedDict
 
 from conllu.exceptions import ParseException
-from conllu.models import Metadata, Token
+from conllu.models import Metadata, Token, TokenList
+
+_IdType = T.Union[int, T.Tuple[int, str, int]]
+_FieldParserType = T.Callable[[T.List[str], int], T.Any]
+_MetadataParserType = T.Callable[[str, T.Optional[str]], T.Any]
 
 DEFAULT_FIELDS = ('id', 'form', 'lemma', 'upos', 'xpos', 'feats', 'head', 'deprel', 'deps', 'misc')
-DEFAULT_FIELD_PARSERS = {
+DEFAULT_FIELD_PARSERS: T.Dict[str, _FieldParserType] = {
     "id": lambda line, i: parse_id_value(line[i]),
     "xpos": lambda line, i: parse_nullable_value(line[i]),
     "feats": lambda line, i: parse_dict_value(line[i]),
@@ -15,12 +20,14 @@ DEFAULT_FIELD_PARSERS = {
     "deps": lambda line, i: parse_paired_list_value(line[i]),
     "misc": lambda line, i: parse_dict_value(line[i]),
 }
-DEFAULT_METADATA_PARSERS = {
+DEFAULT_METADATA_PARSERS: T.Dict[str, _MetadataParserType] = {
     "newpar": lambda key, value: (key, value),
     "newdoc": lambda key, value: (key, value),
 }
 
-def parse_conllu_plus_fields(in_file, metadata_parsers=None):
+def parse_conllu_plus_fields(in_file: T.TextIO,
+                             metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None
+                             ) -> T.Optional[T.Sequence[str]]:
     pos = in_file.tell()
 
     # Get first line
@@ -34,7 +41,7 @@ def parse_conllu_plus_fields(in_file, metadata_parsers=None):
     in_file.seek(pos)
 
     if not first_line.startswith("#"):
-        return
+        return None
 
     _, metadata = parse_token_and_metadata(first_line, metadata_parsers=metadata_parsers)
 
@@ -44,8 +51,8 @@ def parse_conllu_plus_fields(in_file, metadata_parsers=None):
 
     return fields
 
-def parse_sentences(in_file):
-    buf = []
+def parse_sentences(in_file: T.TextIO) -> T.Iterator[str]:
+    buf: T.List[str] = []
     for line in in_file:
         if line == "\n":
             if not buf:
@@ -57,7 +64,10 @@ def parse_sentences(in_file):
     if buf:
         yield "".join(buf).rstrip()
 
-def parse_token_and_metadata(data, fields=None, field_parsers=None, metadata_parsers=None):
+def parse_token_and_metadata(data: str, fields: T.Optional[T.Sequence[str]] = None,
+                             field_parsers: T.Optional[T.Dict[str, _FieldParserType]] = None,
+                             metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None
+                             ) -> T.Tuple[T.List[Token], Metadata]:
     if not data:
         raise ParseException("Can't create TokenList, no data sent to constructor.")
 
@@ -88,7 +98,9 @@ def parse_token_and_metadata(data, fields=None, field_parsers=None, metadata_par
 
     return tokens, metadata
 
-def parse_line(line, fields, field_parsers=None):
+def parse_line(line: str,
+               fields: T.Sequence[str], field_parsers: T.Optional[T.Dict[str, _FieldParserType]] = None
+               ) -> Token:
     # Be backwards compatible if people called parse_line without field_parsers before
     field_parsers = field_parsers or DEFAULT_FIELD_PARSERS
 
@@ -103,32 +115,34 @@ def parse_line(line, fields, field_parsers=None):
     if "upos" not in field_parsers and "upostag" in field_parsers:
         field_parsers["upos"] = field_parsers["upostag"]
 
-    line = re.split(r"\t| {2,}", line)
+    line_split = re.split(r"\t| {2,}", line)
 
-    if len(line) == 1:
+    if len(line_split) == 1:
         raise ParseException("Invalid line format, line must contain either tabs or two spaces.")
 
     data = Token()
 
     for i, field in enumerate(fields):
         # Allow parsing CoNNL-U files with fewer columns
-        if i >= len(line):
+        if i >= len(line_split):
             break
 
         if field in field_parsers:
             try:
-                value = field_parsers[field](line, i)
+                value = field_parsers[field](line_split, i)
             except ParseException as e:
                 raise ParseException("Failed parsing field '{}': ".format(field) + str(e))
 
         else:
-            value = line[i]
+            value = line_split[i]
 
         data[str(field)] = value
 
     return data
 
-def parse_comment_line(line, metadata_parsers=None):
+def parse_comment_line(line: str,
+                       metadata_parsers: T.Optional[T.Dict[str, _MetadataParserType]] = None
+                       ) -> T.List[T.Tuple[str, T.Optional[str]]]:
     line = line.strip()
 
     if line[0] != '#':
@@ -162,17 +176,17 @@ def parse_comment_line(line, metadata_parsers=None):
 
     return [(str(key), value)]
 
-def parse_pair_value(value):
+def parse_pair_value(value: str) -> T.Tuple[str, T.Optional[str]]:
     key_maybe_value = value.split('=', 1)
     key = key_maybe_value[0].strip()
-    value = None if len(key_maybe_value) == 1 else key_maybe_value[1].strip()
+    value_new = None if len(key_maybe_value) == 1 else key_maybe_value[1].strip()
 
-    return key, value
+    return key, value_new
 
 
 INTEGER = re.compile(r"0|(\-?[1-9][0-9]*)")
 
-def parse_int_value(value):
+def parse_int_value(value: str) -> T.Optional[int]:
     if value == '_':
         return None
 
@@ -186,7 +200,7 @@ ID_SINGLE = re.compile(r"(?:0|[1-9][0-9]*)")
 ID_RANGE = re.compile(r"[1-9][0-9]*\-[1-9][0-9]*")
 ID_DOT_ID = re.compile(r"[0-9][0-9]*\.[1-9][0-9]*")
 
-def parse_id_value(value):
+def parse_id_value(value: str) -> T.Optional[_IdType]:
     if not value or value == '_':
         return None
 
@@ -194,8 +208,8 @@ def parse_id_value(value):
         return int(value)
 
     elif re.fullmatch(ID_RANGE, value):
-        from_, to = value.split("-")
-        from_, to = int(from_), int(to)
+        from_str, to_str = value.split("-")
+        from_, to = int(from_str), int(to_str)
         if to > from_:
             return (int(from_), "-", int(to))
 
@@ -209,7 +223,7 @@ ANY_ID = re.compile(ID_SINGLE.pattern + "|" + ID_RANGE.pattern + "|" + ID_DOT_ID
 DEPS_RE = re.compile("(" + ANY_ID.pattern + r")(:[^\d:_\-|][^:|]*)+")
 MULTI_DEPS_PATTERN = re.compile(r"{}(\|{})*".format(DEPS_RE.pattern, DEPS_RE.pattern))
 
-def parse_paired_list_value(value):
+def parse_paired_list_value(value: str) -> T.Union[T.Optional[str], T.List[T.Tuple[str, T.Optional[_IdType]]]]:
     if re.fullmatch(MULTI_DEPS_PATTERN, value):
         return [
             (part.split(":", 1)[1], parse_id_value(part.split(":")[0]))
@@ -218,7 +232,7 @@ def parse_paired_list_value(value):
 
     return parse_nullable_value(value)
 
-def parse_dict_value(value):
+def parse_dict_value(value: str) -> T.Optional[T.Dict[str, T.Optional[str]]]:
     if parse_nullable_value(value) is None:
         return None
 
@@ -227,22 +241,20 @@ def parse_dict_value(value):
         for part in value.split("|") if parse_nullable_value(part.split("=")[0]) is not None
     ])
 
-def parse_nullable_value(value):
+def parse_nullable_value(value: str) -> T.Optional[str]:
     if not value or value == "_":
         return None
 
     return value
 
 # DEPRECATED: Mantain old paths until next major version
-
-def serialize(*args, **kwargs):
+def serialize(tokenlist: TokenList) -> str:
     from conllu.serializer import serialize as new_serialize
-    return new_serialize(*args, **kwargs)
+    return new_serialize(tokenlist)
 
-def serialize_field(*args, **kwargs):
+def serialize_field(field: T.Any) -> str:
     from conllu.serializer import serialize_field as new_serialize_field
-    return new_serialize_field(*args, **kwargs)
+    return new_serialize_field(field)
 
-def head_to_token(*args, **kwargs):
-    from conllu.models import TokenList
-    return TokenList.head_to_token(*args, **kwargs)
+def head_to_token(sentence: TokenList) -> T.Dict[int, T.List[Token]]:
+    return TokenList.head_to_token(sentence)
diff --git a/conllu/py.typed b/conllu/py.typed
new file mode 100644
index 0000000..e69de29
diff --git a/conllu/serializer.py b/conllu/serializer.py
index 4bb1b1e..56dfa2d 100644
--- a/conllu/serializer.py
+++ b/conllu/serializer.py
@@ -1,11 +1,15 @@
 from __future__ import unicode_literals
 
+import typing as T
 from collections import OrderedDict
 
 from conllu.exceptions import ParseException
 
+if T.TYPE_CHECKING:
+    from conllu.models import TokenList
 
-def serialize_field(field):
+
+def serialize_field(field: T.Any) -> str:
     if field is None:
         return '_'
 
@@ -29,7 +33,7 @@ def serialize_field(field):
 
     return "{}".format(field)
 
-def serialize(tokenlist):
+def serialize(tokenlist: 'TokenList') -> str:
     lines = []
 
     if tokenlist.metadata:
diff --git a/setup.cfg b/setup.cfg
index a43894f..86c562c 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -16,6 +16,8 @@ fail_under = 100
 exclude_lines =
 	pragma: no cover
 	def __repr__
+	if T\.TYPE_CHECKING
+	TT\.Protocol
 
 [coverage:run]
 branch = True
diff --git a/setup.py b/setup.py
index f3a4d91..b4030ff 100644
--- a/setup.py
+++ b/setup.py
@@ -1,13 +1,16 @@
 # -*- coding: utf-8 -*-
 import os
 
-from setuptools import setup
+from setuptools import setup  # type: ignore
 
 VERSION = '4.0'
 
 setup(
     name='conllu',
     packages=["conllu"],
+    package_data={
+        "": ["py.typed"]
+    },
     version=VERSION,
     description='CoNLL-U Parser parses a CoNLL-U formatted string into a nested python dictionary',
     long_description=open(os.path.join(os.path.dirname(__file__), 'README.md')).read(),
@@ -15,7 +18,6 @@ setup(
     author=u'Emil Stenström',
     author_email='em@kth.se',
     url='https://github.com/EmilStenstrom/conllu/',
-    install_requires=[],
     keywords=['conllu', 'conll', 'conll-u', 'parser', 'nlp'],
     classifiers=[
         "Programming Language :: Python",
diff --git a/tox.ini b/tox.ini
index e276253..8624bd8 100644
--- a/tox.ini
+++ b/tox.ini
@@ -6,7 +6,7 @@
 [tox]
 envlist =
   py36,py37,py38
-  coverage,flake8,isort
+  coverage,flake8,isort,mypy
 
 [testenv]
 whitelist_externals =
@@ -48,11 +48,16 @@ commands =
   coverage report -m --fail-under=100
   {[cleanup]commands}
 
+[testenv:mypy]
+# Note: Settings for coverage exists in the setup.cfg file
+changedir = {toxinidir}
+deps = mypy
+commands =
+  mypy .
+  {[cleanup]commands}
+
 [travis]
 python =
-  2.7: py27
-  3.4: py34
-  3.5: py35
   3.6: py36
   3.7: py37
-  3.8: py38, coverage, flake8, isort
+  3.8: py38, coverage, flake8, isort, mypy


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
index 36edd0f..f6bab63 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -4,7 +4,7 @@ from __future__ import unicode_literals
 import unittest
 from textwrap import dedent
 
-from conllu.models import Token, TokenList, TokenTree
+from conllu.models import Metadata, Token, TokenList, TokenTree
 from conllu.parser import ParseException, serialize
 from tests.helpers import capture_print
 
@@ -61,7 +61,7 @@ class TestTokenList(unittest.TestCase):
         tokenlist = TokenList([{"id": 1}, {"id": 2}, {"id": 3}], {"meta": "data"})
         tokenlist.clear()
         self.assertEqual(len(tokenlist.tokens), 0)
-        self.assertEqual(tokenlist.metadata, None)
+        self.assertEqual(tokenlist.metadata, Metadata())
 
     def test_copy(self):
         tokenlist1 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}], {"meta": "data"})
@@ -69,13 +69,14 @@ class TestTokenList(unittest.TestCase):
         self.assertIsNot(tokenlist1, tokenlist2)
         self.assertEqual(tokenlist1, tokenlist2)
 
-    def test_extend(self):
+    def test_extend_tokenlist_no_metadata_with_list(self):
         tokenlist1 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}])
         tokenlist2 = [{"id": 4}, {"id": 5}, {"id": 6}]
         tokenlist1.extend(tokenlist2)
         tokenlist3 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}, {"id": 5}, {"id": 6}])
         self.assertEqual(tokenlist1, tokenlist3)
 
+    def test_extend_tokenlist_and_merge_metadata(self):
         tokenlist4 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}], {"meta1": "data1"})
         tokenlist5 = TokenList([{"id": 4}, {"id": 5}, {"id": 6}], {"meta2": "data2"})
         tokenlist4.extend(tokenlist5)
@@ -83,19 +84,6 @@ class TestTokenList(unittest.TestCase):
                                {"meta1": "data1", "meta2": "data2"})
         self.assertEqual(tokenlist4, tokenlist6)
 
-        tokenlist7 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}], "abc")
-        tokenlist8 = TokenList([{"id": 4}, {"id": 5}, {"id": 6}], "de")
-        tokenlist7.extend(tokenlist8)
-        tokenlist9 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}, {"id": 5}, {"id": 6}], "abcde")
-        self.assertEqual(tokenlist7, tokenlist9)
-
-        tokenlist7 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}], "abc")
-        tokenlist8 = TokenList([{"id": 4}, {"id": 5}, {"id": 6}], {"meta2": "data2"})
-        tokenlist7.extend(tokenlist8)
-        tokenlist9 = TokenList([{"id": 1}, {"id": 2}, {"id": 3}, {"id": 4}, {"id": 5}, {"id": 6}],
-                               ["abc", {"meta2": "data2"}])
-        self.assertEqual(tokenlist7, tokenlist9)
-
     def test_tokens(self):
         tokenlist = TokenList([{"id": 1}, {"id": 2}, {"id": 3}])
         self.assertEqual(tokenlist.tokens, [{"id": 1}, {"id": 2}, {"id": 3}])
@@ -275,6 +263,9 @@ class TestTokenTree(unittest.TestCase):
         tokentree2 = TokenTree(token={"id": 1}, children=[])
         self.assertNotEqual(tokentree1, tokentree2)
 
+        # Not equal with a plain dict
+        self.assertNotEqual(tokentree2, {"id": 1})
+
         tokentree2.metadata = metadata
         self.assertNotEqual(tokentree1, tokentree2)