Improve Screener & docs

doc/source/reference/index.rst

@@ -20,7 +20,8 @@ The following are the publicly available classes, and functions exposed by the `
 - :attr:`Sector <yfinance.Sector>`: Domain class for accessing sector information.
 - :attr:`Industry <yfinance.Industry>`: Domain class for accessing industry information.
 - :attr:`download <yfinance.download>`: Function to download market data for multiple tickers.
-- :attr:`EquityQuery <yfinance.EquityQuery>`: Class to build equity market query.
+- :attr:`EquityOperation <yfinance.EquityOperation>`: Class to build equity market operation.
+- :attr:`Query <yfinance.Query>`: Class to build query.
 - :attr:`Screener <yfinance.Screener>`: Class to screen the market using defined query.
 - :attr:`enable_debug_mode <yfinance.enable_debug_mode>`: Function to enable debug mode for logging.
 - :attr:`set_tz_cache_location <yfinance.set_tz_cache_location>`: Function to set the timezone cache location.
@@ -37,6 +38,7 @@ The following are the publicly available classes, and functions exposed by the `
    yfinance.marketsummary
    yfinance.search
    yfinance.sector_industry
+   yfinance.screener
    yfinance.functions
 
    yfinance.funds_data

doc/source/reference/yfinance.functions.rst

@@ -13,25 +13,6 @@ The `download` function allows you to retrieve market data for multiple tickers
 
    download
 
-Query Market Data
-~~~~~~~~~~~~~~~~~~~~~
-The `Sector` and `Industry` modules allow you to access the sector and industry information.
-
-.. autosummary:: 
-   :toctree: api/
-
-   EquityQuery
-   Screener
-
-.. seealso::
-   :attr:`EquityQuery.valid_operand_fields <yfinance.EquityQuery.valid_operand_fields>`
-      supported operand values for query
-   :attr:`EquityQuery.valid_eq_operand_map <yfinance.EquityQuery.valid_eq_operand_map>`
-      supported `EQ query operand parameters`
-   :attr:`Screener.predefined_bodies <yfinance.Screener.predefined_bodies>`
-      supported predefined screens
-
-
 Enable Debug Mode
 ~~~~~~~~~~~~~~~~~
 Enables logging of debug information for the `yfinance` package.

doc/source/reference/yfinance.screener.rst

@@ -0,0 +1,27 @@
+=========================
+Screener & Query
+=========================
+
+.. currentmodule:: yfinance
+
+Query Market Data
+~~~~~~~~~~~~~~~~~~~~~
+The `Sector` and `Industry` modules allow you to access the sector and industry information.
+
+.. autosummary:: 
+   :toctree: api/
+
+   EquityQuery
+   FundQuery
+   screen
+
+.. seealso::
+   :attr:`EquityQuery.valid_fields <yfinance.EquityQuery.valid_fields>`
+      supported operand values for query
+   :attr:`EquityQuery.valid_values <yfinance.EquityQuery.valid_values>`
+      supported `EQ query operand parameters`
+   :attr:`EquityQuery.valid_fields <yfinance.EquityQuery.valid_fields>`
+      supported operand values for query
+   :attr:`EquityQuery.valid_values <yfinance.EquityQuery.valid_values>`
+      supported `EQ query operand parameters`
+

tests/test_screener.py

@@ -1,153 +1,38 @@
 import unittest
 from unittest.mock import patch, MagicMock
-from yfinance.const import PREDEFINED_SCREENER_BODY_MAP
-from yfinance.screener.screener import Screener
-from yfinance.screener.screener_query import EquityQuery
+from yfinance.screener.screener import screen
+from yfinance.screener.query import EquityQuery
 
 
 class TestScreener(unittest.TestCase):
 
     @classmethod
     def setUpClass(self):
-        self.screener = Screener()
         self.query = EquityQuery('gt',['eodprice',3])
-
-    def test_set_default_body(self):
-        result = self.screener.set_default_body(self.query)
-
-        self.assertEqual(self.screener.body['offset'], 0)
-        self.assertEqual(self.screener.body['size'], 100)
-        self.assertEqual(self.screener.body['sortField'], 'ticker')
-        self.assertEqual(self.screener.body['sortType'], 'desc')
-        self.assertEqual(self.screener.body['quoteType'], 'equity')
-        self.assertEqual(self.screener.body['query'], self.query.to_dict())
-        self.assertEqual(self.screener.body['userId'], '')
-        self.assertEqual(self.screener.body['userIdType'], 'guid')
-        self.assertEqual(self.screener, result)
-
-    def test_set_predefined_body(self):
-        k = 'most_actives'
-        result = self.screener.set_predefined_body(k)
-        self.assertEqual(self.screener.body, PREDEFINED_SCREENER_BODY_MAP[k])
-        self.assertEqual(self.screener, result)
-
-    def test_set_predefined_body_invalid_key(self):
-        with self.assertRaises(ValueError):
-            self.screener.set_predefined_body('invalid_key')
-
-    def test_set_body(self):
-        body = {
-            "offset": 0,
-            "size": 100,
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity",
-            "query": self.query.to_dict(),
-            "userId": "",
-            "userIdType": "guid"
-        }
-        result = self.screener.set_body(body)
-
-        self.assertEqual(self.screener.body, body)
-        self.assertEqual(self.screener, result)
-
-    def test_set_body_missing_keys(self):
-        body = {
-            "offset": 0,
-            "size": 100,
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity"
-        }
-        with self.assertRaises(ValueError):
-            self.screener.set_body(body)
-
-    def test_set_body_extra_keys(self):
-        body = {
-            "offset": 0,
-            "size": 100,
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity",
-            "query": self.query.to_dict(),
-            "userId": "",
-            "userIdType": "guid",
-            "extraKey": "extraValue"
-        }
-        with self.assertRaises(ValueError):
-            self.screener.set_body(body)
-
-    def test_patch_body(self):
-        initial_body = {
-            "offset": 0,
-            "size": 100,
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity",
-            "query": self.query.to_dict(),
-            "userId": "",
-            "userIdType": "guid"
-        }
-        self.screener.set_body(initial_body)
-        patch_values = {"size": 50}
-        result = self.screener.patch_body(patch_values)
-
-        self.assertEqual(self.screener.body['size'], 50)
-        self.assertEqual(self.screener.body['query'], self.query.to_dict())
-        self.assertEqual(self.screener, result)
-
-    def test_patch_body_extra_keys(self):
-        initial_body = {
-            "offset": 0,
-            "size": 100,
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity",
-            "query": self.query.to_dict(),
-            "userId": "",
-            "userIdType": "guid"
-        }
-        self.screener.set_body(initial_body)
-        patch_values = {"extraKey": "extraValue"}
-        with self.assertRaises(ValueError):
-            self.screener.patch_body(patch_values)
+        self.predefined = 'aggressive_small_caps'
 
     @patch('yfinance.screener.screener.YfData.post')
     def test_set_large_size_in_body(self, mock_post):
-        body = {
-            "offset": 0,
-            "size": 251, # yahoo limits at 250
-            "sortField": "ticker",
-            "sortType": "desc",
-            "quoteType": "equity",
-            "query": self.query.to_dict(),
-            "userId": "",
-            "userIdType": "guid"
-        }
-
         with self.assertRaises(ValueError):
-            self.screener.set_body(body).response            
+            screen(self.query, size=251)
 
-    @patch('yfinance.screener.screener.YfData.post')
-    def test_fetch(self, mock_post):
+    @patch('yfinance.data.YfData.post')
+    def test_fetch_query(self, mock_post):
         mock_response = MagicMock()
-        mock_response.json.return_value = {'finance': {'result': [{}]}}
+        mock_response.json.return_value = {'finance': {'result': [{'key': 'value'}]}}
         mock_post.return_value = mock_response
 
-        self.screener.set_default_body(self.query)
-        response = self.screener._fetch()
-
-        self.assertEqual(response, {'finance': {'result': [{}]}})
+        response = screen(self.query)
+        self.assertEqual(response, {'key': 'value'})
 
-    @patch('yfinance.screener.screener.YfData.post')
-    def test_fetch_and_parse(self, mock_post):
+    @patch('yfinance.data.YfData.get')
+    def test_fetch_predefined(self, mock_get):
         mock_response = MagicMock()
         mock_response.json.return_value = {'finance': {'result': [{'key': 'value'}]}}
-        mock_post.return_value = mock_response
+        mock_get.return_value = mock_response
 
-        self.screener.set_default_body(self.query)
-        self.screener._fetch_and_parse()
-        self.assertEqual(self.screener.response, {'key': 'value'})
+        response = screen(self.predefined)
+        self.assertEqual(response, {'key': 'value'})
 
 if __name__ == '__main__':
     unittest.main()

yfinance/__init__.py

@@ -28,15 +28,17 @@
 from .cache import set_tz_cache_location
 from .domain.sector import Sector
 from .domain.industry import Industry
-from .screener.screener import Screener
-from .screener.screener_query import EquityQuery
 from .domain.market import Market
 
+from .screener.query import EquityQuery, FundQuery
+from .screener.screener import screen, PREDEFINED_SCREENER_QUERIES
+
 __version__ = version.version
 __author__ = "Ran Aroussi"
 
 import warnings
 warnings.filterwarnings('default', category=DeprecationWarning, module='^yfinance')
 
-__all__ = ['download', 'Market', 'Search', 'Ticker', 'Tickers', 'enable_debug_mode', 'set_tz_cache_location', 'Sector',
-           'Industry', 'EquityQuery', 'Screener']
+__all__ = ['download', 'Market', 'Search', 'Ticker', 'Tickers', 'enable_debug_mode', 'set_tz_cache_location', 'Sector', 'Industry']
+# screener stuff:
+__all__ += ['EquityQuery', 'FundQuery', 'screen', 'PREDEFINED_SCREENER_QUERIES']