pixeltable 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/tests/conftest.py

@@ -1,119 +1,175 @@
-import numpy as np
-import pandas as pd
-import datetime
+import json
+import logging
+import os
+import pathlib
+from typing import List
 
+import numpy as np
 import pytest
 
-import pixeltable as pt
+import pixeltable as pxt
 import pixeltable.catalog as catalog
-from pixeltable.type_system import \
-    StringType, IntType, FloatType, BoolType, TimestampType, ImageType, JsonType
-from pixeltable.tests.utils import read_data_file, make_tbl, create_table_data
+from pixeltable import exprs
+from pixeltable import functions as ptf
+from pixeltable.exprs import RELATIVE_PATH_ROOT as R
+from pixeltable.metadata import SystemInfo, create_system_info
+from pixeltable.metadata.schema import TableSchemaVersion, TableVersion, Table, Function, Dir
+from pixeltable.tests.utils import read_data_file, create_test_tbl, create_all_datatypes_tbl, skip_test_if_not_installed
+from pixeltable.type_system import StringType, ImageType, FloatType
 
 
 @pytest.fixture(scope='session')
-def init_db(tmp_path_factory) -> None:
+def init_env(tmp_path_factory) -> None:
     from pixeltable.env import Env
+    # set the relevant env vars for Client() to connect to the test db
+
+    shared_home = pathlib.Path(os.environ.get('PIXELTABLE_HOME', '~/.pixeltable')).expanduser()
+    home_dir = str(tmp_path_factory.mktemp('base') / '.pixeltable')
+    os.environ['PIXELTABLE_HOME'] = home_dir
+    os.environ['PIXELTABLE_CONFIG'] = str(shared_home / 'config.yaml')
+    test_db = 'test'
+    os.environ['PIXELTABLE_DB'] = test_db
+    os.environ['PIXELTABLE_PGDATA'] = str(shared_home / 'pgdata')
+
     # this also runs create_all()
-    db_name = 'test'
-    Env.get().set_up(tmp_path_factory.mktemp('base'), db_name=db_name, echo=True)
+    Env.get().set_up(echo=True)
     yield
     # leave db in place for debugging purposes
 
+@pytest.fixture(scope='function')
+def test_client(init_env) -> pxt.Client:
+    # Clean the DB *before* instantiating a client object. This is because some tests
+    # (such as test_migration.py) may leave the DB in a broken state, from which the
+    # client is uninstantiable.
+    clean_db()
+    cl = pxt.Client(reload=True)
+    cl.logging(level=logging.DEBUG, to_stdout=True)
+    yield cl
+
+
+def clean_db(restore_tables: bool = True) -> None:
+    from pixeltable.env import Env
+    # The logic from Client.reset_catalog() has been moved here, so that it
+    # does not rely on instantiating a Client object. As before, UUID-named data tables will
+    # not be cleaned. If in the future it is desirable to clean out data tables as well,
+    # the commented lines may be used to drop ALL tables from the test db.
+    # sql_md = declarative_base().metadata
+    # sql_md.reflect(Env.get().engine)
+    # sql_md.drop_all(bind=Env.get().engine)
+    engine = Env.get().engine
+    SystemInfo.__table__.drop(engine, checkfirst=True)
+    TableSchemaVersion.__table__.drop(engine, checkfirst=True)
+    TableVersion.__table__.drop(engine, checkfirst=True)
+    Table.__table__.drop(engine, checkfirst=True)
+    Function.__table__.drop(engine, checkfirst=True)
+    Dir.__table__.drop(engine, checkfirst=True)
+    if restore_tables:
+        Dir.__table__.create(engine)
+        Function.__table__.create(engine)
+        Table.__table__.create(engine)
+        TableVersion.__table__.create(engine)
+        TableSchemaVersion.__table__.create(engine)
+        SystemInfo.__table__.create(engine)
+        create_system_info(engine)
+
 
 @pytest.fixture(scope='function')
-def test_db(init_db: None) -> pt.Db:
-    cl = pt.Client()
-    db = cl.create_db(f'test')
-    yield db
-    cl.drop_db(db.name, force=True)
+def test_tbl(test_client: pxt.Client) -> catalog.Table:
+    return create_test_tbl(test_client)
 
+# @pytest.fixture(scope='function')
+# def test_stored_fn(test_client: pt.Client) -> pt.Function:
+#     @pt.udf(return_type=pt.IntType(), param_types=[pt.IntType()])
+#     def test_fn(x):
+#         return x + 1
+#     test_client.create_function('test_fn', test_fn)
+#     return test_fn
 
 @pytest.fixture(scope='function')
-def test_tbl(test_db: pt.Db) -> catalog.Table:
-    cols = [
-        catalog.Column('c1', StringType(), nullable=False),
-        catalog.Column('c2', IntType(), nullable=False),
-        catalog.Column('c3', FloatType(), nullable=False),
-        catalog.Column('c4', BoolType(), nullable=False),
-        catalog.Column('c5', TimestampType(), nullable=False),
-        catalog.Column('c6', JsonType(), nullable=False),
-        catalog.Column('c7', JsonType(), nullable=False),
+def test_tbl_exprs(test_tbl: catalog.Table) -> List[exprs.Expr]:
+#def test_tbl_exprs(test_tbl: catalog.Table, test_stored_fn: pt.Function) -> List[exprs.Expr]:
+
+    t = test_tbl
+    return [
+        t.c1,
+        t.c7['*'].f1,
+        exprs.Literal('test'),
+        exprs.InlineDict({
+            'a': t.c1, 'b': t.c6.f1, 'c': 17,
+            'd': exprs.InlineDict({'e': t.c2}),
+            'f': exprs.InlineArray((t.c3, t.c3))
+        }),
+        exprs.InlineArray([[t.c2, t.c2], [t.c2, t.c2]]),
+        t.c2 > 5,
+        t.c2 == None,
+        ~(t.c2 > 5),
+        (t.c2 > 5) & (t.c1 == 'test'),
+        (t.c2 > 5) | (t.c1 == 'test'),
+        t.c7['*'].f5 >> [R[3], R[2], R[1], R[0]],
+        t.c8[0, 1:],
+        t.c2.astype(FloatType()),
+        (t.c2 + 1).astype(FloatType()),
+        t.c2.apply(str),
+        (t.c2 + 1).apply(str),
+        t.c3.apply(str),
+        t.c4.apply(str),
+        t.c5.apply(str),
+        t.c6.apply(str),
+        t.c1.apply(json.loads),
+        t.c8.errortype,
+        t.c8.errormsg,
+        ptf.sum(t.c2, group_by=t.c4, order_by=t.c3),
+        #test_stored_fn(t.c2),
     ]
-    t = test_db.create_table('test__tbl', cols)
-
-    num_rows = 100
-    d1 = {
-        'f1': 'test string 1',
-        'f2': 1,
-        'f3': 1.0,
-        'f4': True,
-        'f5': [1.0, 2.0, 3.0, 4.0],
-        'f6': {
-            'f7': 'test string 2',
-            'f8': [1.0, 2.0, 3.0, 4.0],
-        },
-    }
-    d2 = [d1, d1]
-
-    c1_data = [f'test string {i}' for i in range(num_rows)]
-    c2_data = [i for i in range(num_rows)]
-    c3_data = [float(i) for i in range(num_rows)]
-    c4_data = [bool(i % 2) for i in range(num_rows)]
-    c5_data = [datetime.datetime.now()] * num_rows
-    c6_data = []
-    for i in range(num_rows):
-        d = {
-            'f1': f'test string {i}',
-            'f2': i,
-            'f3': float(i),
-            'f4': bool(i % 2),
-            'f5': [1.0, 2.0, 3.0, 4.0],
-            'f6': {
-                'f7': 'test string 2',
-                'f8': [1.0, 2.0, 3.0, 4.0],
-            },
-        }
-        c6_data.append(d)
-
-    c7_data = [d2] * num_rows
-    data = {'c1': c1_data, 'c2': c2_data, 'c3': c3_data, 'c4': c4_data, 'c5': c5_data, 'c6': c6_data, 'c7': c7_data}
-    pd_df = pd.DataFrame(data=data)
-    t.insert_pandas(pd_df)
-    return t
 
+@pytest.fixture(scope='function')
+def all_datatypes_tbl(test_client: pxt.Client) -> catalog.Table:
+    return create_all_datatypes_tbl(test_client)
 
 @pytest.fixture(scope='function')
-def img_tbl(test_db: pt.Db) -> catalog.Table:
-    cols = [
-        catalog.Column('img', ImageType(), nullable=False, indexed=False),
-        catalog.Column('category', StringType(), nullable=False),
-        catalog.Column('split', StringType(), nullable=False),
-    ]
+def img_tbl(test_client: pxt.Client) -> catalog.Table:
+    schema = {
+        'img': ImageType(nullable=False),
+        'category': StringType(nullable=False),
+        'split': StringType(nullable=False),
+    }
     # this table is not indexed in order to avoid the cost of computing embeddings
-    tbl = test_db.create_table('test_img_tbl', cols)
-    df = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
-    tbl.insert_pandas(df)
+    tbl = test_client.create_table('test_img_tbl', schema)
+    rows = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
+    tbl.insert(rows)
     return tbl
 
+@pytest.fixture(scope='function')
+def img_tbl_exprs(img_tbl: catalog.Table) -> List[exprs.Expr]:
+    img_t = img_tbl
+    return [
+        img_t.img.width,
+        img_t.img.rotate(90),
+        # we're using a list here, not a tuple; the latter turns into a list during the back/forth conversion
+        img_t.img.rotate(90).resize([224, 224]),
+        img_t.img.fileurl,
+        img_t.img.localpath,
+    ]
 
 # TODO: why does this not work with a session scope? (some user tables don't get created with create_all())
 #@pytest.fixture(scope='session')
-#def indexed_img_tbl(init_db: None) -> catalog.Table:
+#def indexed_img_tbl(init_env: None) -> catalog.Table:
 #    cl = pt.Client()
 #    db = cl.create_db('test_indexed')
 @pytest.fixture(scope='function')
-def indexed_img_tbl(test_db: pt.Db) -> catalog.Table:
-    db = test_db
-    cols = [
-        catalog.Column('img', ImageType(), nullable=False, indexed=True),
-        catalog.Column('category', StringType(), nullable=False),
-        catalog.Column('split', StringType(), nullable=False),
-    ]
-    tbl = db.create_table('test_indexed_img_tbl', cols)
-    df = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
-    # select rows randomly in the hope of getting a good sample of the available categories
+def indexed_img_tbl(test_client: pxt.Client) -> catalog.Table:
+    skip_test_if_not_installed('nos')
+    cl = test_client
+    schema = {
+        'img': { 'type': ImageType(nullable=False), 'indexed': True },
+        'category': StringType(nullable=False),
+        'split': StringType(nullable=False),
+    }
+    tbl = cl.create_table('test_indexed_img_tbl', schema)
+    rows = read_data_file('imagenette2-160', 'manifest.csv', ['img'])
+    # select output_rows randomly in the hope of getting a good sample of the available categories
     rng = np.random.default_rng(17)
-    idxs = rng.choice(np.arange(len(df)), size=40, replace=False)
-    tbl.insert_pandas(df.iloc[idxs])
+    idxs = rng.choice(np.arange(len(rows)), size=40, replace=False)
+    rows = [rows[i] for i in idxs]
+    tbl.insert(rows)
     return tbl

pixeltable/tests/test_audio.py

@@ -0,0 +1,65 @@
+from typing import Optional
+
+import av
+
+import pixeltable as pxt
+import pixeltable.env as env
+from pixeltable.tests.utils import get_video_files, get_audio_files
+from pixeltable.type_system import VideoType, AudioType
+from pixeltable.utils.media_store import MediaStore
+
+
+class TestAudio:
+    def check_audio_params(self, path: str, format: Optional[str] = None, codec: Optional[str] = None):
+        with av.open(path) as container:
+            audio_stream = container.streams.audio[0]
+            if format is not None:
+                assert format == container.format.name
+            if codec is not None:
+                assert codec == audio_stream.codec_context.codec.name
+
+    def test_basic(self, test_client: pxt.Client) -> None:
+        audio_filepaths = get_audio_files()
+        cl = test_client
+        audio_t = cl.create_table('audio', {'audio_file': AudioType()})
+        status = audio_t.insert({'audio_file': p} for p in audio_filepaths)
+        assert status.num_rows == len(audio_filepaths)
+        assert status.num_excs == 0
+        paths = audio_t.select(output=audio_t.audio_file.localpath).collect()['output']
+        assert set(paths) == set(audio_filepaths)
+
+    def test_extract(self, test_client: pxt.Client) -> None:
+        video_filepaths = get_video_files()
+        cl = test_client
+        video_t = cl.create_table('videos', {'video': VideoType()})
+        from pixeltable.functions.video import extract_audio
+        video_t.add_column(audio=extract_audio(video_t.video))
+
+        # one of the 3 videos doesn't have audio
+        status = video_t.insert({'video': p} for p in video_filepaths)
+        assert status.num_rows == len(video_filepaths)
+        assert status.num_excs == 0
+        assert MediaStore.count(video_t.get_id()) == len(video_filepaths) - 1
+        assert video_t.where(video_t.audio != None).count() == len(video_filepaths) - 1
+        assert env.Env.get().num_tmp_files() == 0
+
+        # make sure everything works with a fresh client
+        cl = pxt.Client()
+        video_t = cl.get_table('videos')
+        assert video_t.where(video_t.audio != None).count() == len(video_filepaths) - 1
+
+        # test generating different formats and codecs
+        paths = video_t.select(output=extract_audio(video_t.video, format='wav', codec='pcm_s16le')).collect()['output']
+        # media files that are created as a part of a query end up in the tmp dir
+        assert env.Env.get().num_tmp_files() == video_t.where(video_t.audio != None).count()
+        for path in [p for p in paths if p is not None]:
+            self.check_audio_params(path, format='wav', codec='pcm_s16le')
+        # higher resolution
+        paths = video_t.select(output=extract_audio(video_t.video, format='wav', codec='pcm_s32le')).collect()['output']
+        for path in [p for p in paths if p is not None]:
+            self.check_audio_params(path, format='wav', codec='pcm_s32le')
+
+        for format in ['mp3', 'flac']:
+            paths = video_t.select(output=extract_audio(video_t.video, format=format)).collect()['output']
+            for path in [p for p in paths if p is not None]:
+                self.check_audio_params(path, format=format)

pixeltable/tests/test_catalog.py

@@ -0,0 +1,27 @@
+from pixeltable.catalog import is_valid_identifier, is_valid_path
+
+class TestCatalog:
+    """Tests for miscellanous catalog functions."""
+    def test_valid_identifier(self) -> None:
+        valid_ids = ['a', 'a1', 'a_1', 'a_']
+        invalid_ids = ['', '_', '__', '_a', '1a', 'a.b', '.a', 'a-b']
+        for valid_id in valid_ids:
+            assert is_valid_identifier(valid_id), valid_ids
+
+        for invalid_id in invalid_ids:
+            assert not is_valid_identifier(invalid_id), invalid_ids
+
+    def test_valid_path(self) -> None:
+        assert is_valid_path('', empty_is_valid=True)
+        assert not is_valid_path('', empty_is_valid=False)
+
+        valid_paths = ['a', 'a_.b_', 'a.b.c', 'a.b.c.d']
+        invalid_paths = ['.', '..', 'a.', '.a', 'a..b']
+        
+        for valid_path in valid_paths:
+            assert is_valid_path(valid_path, empty_is_valid=False), valid_path
+            assert is_valid_path(valid_path, empty_is_valid=True), valid_path
+
+        for invalid_path in invalid_paths:
+            assert not is_valid_path(invalid_path, empty_is_valid=False), invalid_path
+            assert not is_valid_path(invalid_path, empty_is_valid=True), invalid_path

pixeltable/tests/test_client.py

@@ -1,21 +1,21 @@
 import pytest
 
-import pixeltable as pt
-from pixeltable import exceptions as exc
+import pixeltable as pxt
+import pixeltable.exceptions as excs
 
 
 class TestClient:
-    def test_create_db(self, init_db: None) -> None:
-        cl = pt.Client()
-        _ = cl.create_db('test')
-        with pytest.raises(exc.DuplicateNameError):
-            _ = cl.create_db('test')
+    def test_list_functions(self, init_env) -> None:
+        cl = pxt.Client()
+        _ = cl.list_functions()
+        print(_)
 
-        _ = cl.get_db('test')
-        with pytest.raises(exc.UnknownEntityError):
-            _ = cl.get_db('xyz')
-
-        cl.drop_db('test', force=True)
-        with pytest.raises(exc.UnknownEntityError):
-            cl.drop_db('test', force=True)
+    def test_drop_table(self, test_tbl: pxt.Table) -> None:
+        cl = pxt.Client()
+        t = cl.get_table('test_tbl')
+        cl.drop_table('test_tbl')
+        with pytest.raises(excs.Error):
+            _ = cl.get_table('test_tbl')
+        with pytest.raises(excs.Error):
+            _ = t.show(1)