From 55768aefe2eeca38b3d799e9e04f861c95c38572 Mon Sep 17 00:00:00 2001
From: Sugandh <s.sugandh@rocketmail.com>
Date: Mon, 25 Oct 2021 21:58:15 +0200
Subject: [PATCH] Fix getpoly to return only US stocks (#801)

* Fix getpoly to return only US stocks

* Fix addresses #757
* Getpoly will only return stocks listed on US stock exchanges.
* Index/mutual funds will also not be returned by getpoly.

* Remove unneeded else condition

* Add --us_only flag for getpoly to get US stocks only

* Fix black check error
---
 .../stocks/comparison_analysis/ca_controller.py      | 12 +++++++++++-
 .../stocks/comparison_analysis/polygon_model.py      | 11 ++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/gamestonk_terminal/stocks/comparison_analysis/ca_controller.py b/gamestonk_terminal/stocks/comparison_analysis/ca_controller.py
index 40f39be3945..4ca8773824a 100644
--- a/gamestonk_terminal/stocks/comparison_analysis/ca_controller.py
+++ b/gamestonk_terminal/stocks/comparison_analysis/ca_controller.py
@@ -240,12 +240,22 @@ Finviz:
             prog="getpoly",
             description="""Get similar companies from polygon to compare with.""",
         )
+        parser.add_argument(
+            "-u",
+            "--us_only",
+            action="store_true",
+            default=False,
+            dest="us_only",
+            help="Show only stocks from the US stock exchanges",
+        )
 
         try:
             ns_parser = parse_known_args_and_warn(parser, other_args)
             if not ns_parser:
                 return
-            self.similar, self.user = polygon_model.get_similar_companies(self.ticker)
+            self.similar, self.user = polygon_model.get_similar_companies(
+                self.ticker, ns_parser.us_only
+            )
 
             if self.ticker.upper() in self.similar:
                 self.similar.remove(self.ticker.upper())
diff --git a/gamestonk_terminal/stocks/comparison_analysis/polygon_model.py b/gamestonk_terminal/stocks/comparison_analysis/polygon_model.py
index dd84377b258..c2432bced64 100644
--- a/gamestonk_terminal/stocks/comparison_analysis/polygon_model.py
+++ b/gamestonk_terminal/stocks/comparison_analysis/polygon_model.py
@@ -7,7 +7,7 @@ import requests
 from gamestonk_terminal import config_terminal as cfg
 
 
-def get_similar_companies(ticker: str) -> Tuple[List[str], str]:
+def get_similar_companies(ticker: str, us_only: bool) -> Tuple[List[str], str]:
     """Get similar companies from Polygon
 
     Parameters
@@ -29,6 +29,15 @@ def get_similar_companies(ticker: str) -> Tuple[List[str], str]:
     if result.status_code == 200:
         similar = result.json()["similar"]
         source = "Polygon"
+        if us_only:
+            us_similar = []
+            mkw_link = "https://www.marketwatch.com/investing/stock/"
+            for sym in similar:
+                prep_link = mkw_link + sym
+                sent_req = requests.get(prep_link)
+                if prep_link == sent_req.request.url:
+                    us_similar.append(sym)
+                similar = us_similar
     else:
         print(result.json()["error"])
         similar = [""]